YOLOv5 + Albumentations integration (#3882)

* Albumentations integration

* ToGray p=0.01

* print confirmation

* create instance in dataloader init method

* improved version handling

* transform not defined fix

* assert string update

* create check_version()

* add spaces

* update class comment

utils/augmentations.py

@@ -1,15 +1,43 @@
 # YOLOv5 image augmentation functions
 
+import logging
 import random
 
 import cv2
 import math
 import numpy as np
 
-from utils.general import segment2box, resample_segments
+from utils.general import colorstr, segment2box, resample_segments, check_version
 from utils.metrics import bbox_ioa
 
 
+class Albumentations:
+    # YOLOv5 Albumentations class (optional, only used if package is installed)
+    def __init__(self):
+        self.transform = None
+        try:
+            import albumentations as A
+            check_version(A.__version__, '1.0.0')  # version requirement
+
+            self.transform = A.Compose([
+                A.Blur(p=0.1),
+                A.MedianBlur(p=0.1),
+                A.ToGray(p=0.01)],
+                bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))
+
+            logging.info(colorstr('albumentations: ') + ', '.join(f'{x}' for x in self.transform.transforms))
+        except ImportError:  # package not installed, skip
+            pass
+        except Exception as e:
+            logging.info(colorstr('albumentations: ') + f'{e}')
+
+    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'])])
+        return im, labels
+
+
 def augment_hsv(im, hgain=0.5, sgain=0.5, vgain=0.5):
     # HSV color-space augmentation
     if hgain or sgain or vgain: