Skip to content

Instantly share code, notes, and snippets.

@poxyu
Last active October 4, 2018 06:50
Show Gist options
  • Star 3 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save poxyu/799b359d7f87205b7c4288691ab019dc to your computer and use it in GitHub Desktop.
Save poxyu/799b359d7f87205b7c4288691ab019dc to your computer and use it in GitHub Desktop.
Custom targets for albumentations
# import base class first
# or you'll get such error:
# TypeError: super(type, obj): obj must be an instance or subtype of type
from albumentations import DualTransform
# new targets function
def new_tf_targets(self):
return {
'image': self.apply, # do not rename this one
'image2': self.apply, # new
'target': self.apply, # new
'mask': self.apply_to_mask,
'bboxes': self.apply_to_bboxes
}
# redefine targets function
# don't forget to wrap it in property
DualTransform.targets = property(new_tf_targets)
# now import all dependent classes
from albumentations import (
...
)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment