From 429238f7cd38c1041d6f2e57fdfbce585afab4f7 Mon Sep 17 00:00:00 2001
From: UnglvKitDe <unglv@student.kit.edu>
Date: Sat, 23 Jul 2022 22:34:48 +0200
Subject: [PATCH] Fix BGR->RGB Bug in albumentations
 https://github.com/ultralytics/yolov5/issues/8641

---
 utils/augmentations.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/utils/augmentations.py b/utils/augmentations.py
index 3f764c06ae3b..4183cf0e89a2 100644
--- a/utils/augmentations.py
+++ b/utils/augmentations.py
@@ -39,8 +39,10 @@ def __init__(self):
 
     def __call__(self, im, labels, p=1.0):
         if self.transform and random.random() < p:
+            im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
             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'])])
+            im = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)
         return im, labels