Fix BGR->RGB Bug in albumentations #8641 (#8695)

* Fix BGR->RGB Bug in albumentations #8641 * Change transform methode from cv2 to numpy * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Simplify * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update augmentations.py Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Glenn Jocher <[email protected]>
ultralytics · Jul 26, 2022 · 2e1291f · 2e1291f
1 parent b367860
commit 2e1291f
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/utils/augmentations.py b/utils/augmentations.py
@@ -39,8 +39,9 @@ def __init__(self):
 
     def __call__(self, im, labels, p=1.0):
         if self.transform and random.random() < p:
-            new = self.transform(image=im, bboxes=labels[:, 1:], class_labels=labels[:, 0])  # transformed
-            im, labels = new['image'], np.array([[c, *b] for c, b in zip(new['class_labels'], new['bboxes'])])
+            new = self.transform(image=im[..., ::-1], bboxes=labels[:, 1:], class_labels=labels[:, 0])  # transformed
+            im = new['image'][..., ::-1]  # RGB to BGR
+            labels = np.array([[c, *b] for c, b in zip(new['class_labels'], new['bboxes'])])
         return im, labels