From 2e1291fdce26b3cff213e9e7ee8c196fa263b688 Mon Sep 17 00:00:00 2001 From: UnglvKitDe <100289696+UnglvKitDe@users.noreply.github.com> Date: Tue, 26 Jul 2022 13:52:56 +0200 Subject: [PATCH] Fix BGR->RGB Bug in albumentations #8641 (#8695) * Fix BGR->RGB Bug in albumentations https://github.com/ultralytics/yolov5/issues/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 --- utils/augmentations.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/utils/augmentations.py b/utils/augmentations.py index 3f764c0..97506ae 100644 --- a/utils/augmentations.py +++ b/utils/augmentations.py @@ -39,8 +39,9 @@ class Albumentations: 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