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