Disable check_each_transform for rotated labels.

acai66 · Nov 18, 2021 · a4b4dff · a4b4dff
1 parent 9e0771b
commit a4b4dff
Showing 1 changed file with 10 additions and 8 deletions.
diff --git a/utils/augmentations.py b/utils/augmentations.py
@@ -22,14 +22,14 @@ def __init__(self):
             check_version(A.__version__, '1.0.3', hard=True)  # version requirement
 
             self.transform = A.Compose([
-                A.Blur(p=0.01),
-                A.MedianBlur(p=0.01),
-                A.ToGray(p=0.01),
-                A.CLAHE(p=0.01),
-                A.RandomBrightnessContrast(p=0.0),
-                A.RandomGamma(p=0.0),
-                A.ImageCompression(quality_lower=75, p=0.0)],
-                bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))
+                A.Blur(p=0.01, check_each_transform=False),
+                A.MedianBlur(p=0.01, check_each_transform=False),
+                A.ToGray(p=0.01, check_each_transform=False),
+                A.CLAHE(p=0.01, check_each_transform=False),
+                A.RandomBrightnessContrast(p=0.0, check_each_transform=False),
+                A.RandomGamma(p=0.0, check_each_transform=False),
+                A.ImageCompression(quality_lower=75, p=0.0, check_each_transform=False)],
+                bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels'], check_each_transform=False))
 
             LOGGER.info(colorstr('albumentations: ') + ', '.join(f'{x}' for x in self.transform.transforms if x.p))
         except ImportError:  # package not installed, skip
@@ -39,8 +39,10 @@ def __init__(self):
 
     def __call__(self, im, labels, p=1.0):
         if self.transform and random.random() < p:
+            angle_labels = labels[:, 5:6]
             new = self.transform(image=im, bboxes=labels[:, 1:5], class_labels=labels[:, 0])  # transformed
             im, labels = new['image'], np.array([[c, *b] for c, b in zip(new['class_labels'], new['bboxes'])])
+            labels = np.hstack([labels, angle_labels])
         return im, labels