YOLOv5 + Albumentations integration

[glenn-jocher]

I'm super excited to announce our new YOLOv5 🚀 + Albumentations integration!! Now you can train the world's best Vision AI models even better with custom Albumentations automatically applied 😃!

PR #3882 implements this integration, which will automatically apply Albumentations transforms during YOLOv5 training if albumentations>=1.0.3 is installed in your environment.

Get Started

To use albumentations simply pip install -U albumentations and then update the augmentation pipeline as you see fit in the new Albumentations class in yolov5/utils/augmentations.py. Note these Albumentations operations run in addition to the YOLOv5 hyperparameter augmentations, i.e. defined in hyp.scratch.yaml.

Here's an example that applies Blur, MedianBlur and ToGray albumentations in addition to the YOLOv5 hyperparameter augmentations normally applied to your training mosaics :)

class Albumentations:
    # YOLOv5 Albumentations class (optional, used if package is installed)
    def __init__(self):
        self.transform = None
        try:
            import albumentations as A
            check_version(A.__version__, '1.0.3')  # version requirement

            self.transform = A.Compose([
                A.Blur(blur_limit=50, p=0.1),
                A.MedianBlur(blur_limit=51, p=0.1),
                A.ToGray(p=0.3)],
                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

Example Result

Example train_batch0.jpg on COCO128 dataset with Blur, MedianBlur and ToGray. See the YOLOv5 Notebooks to reproduce:

Update

To receive this YOLOv5 update:

• Git – git pull from within your yolov5/ directory or git clone https://github.com/ultralytics/yolov5 again
• PyTorch Hub – Force-reload with model = torch.hub.load('ultralytics/yolov5', 'yolov5s', force_reload=True)
• Notebooks – View updated notebooks
• Docker – sudo docker pull ultralytics/yolov5:latest to update your image

Thank you for your feedback and let us know if this update works for you, and of course feel free to inform us of any other issues you discover or feature requests that come to mind. Happy trainings with YOLOv5 🚀 + Albumentations!

🛠️ PR Summary

_{Made with ❤️ by Ultralytics Actions}

🌟 Summary

Enhanced image augmentation with Albumentations integration in YOLOv5

📊 Key Changes

• 📦 Albumentations library added to requirements.txt as an optional dependency.
• ✨ New Albumentations class created in augmentations.py to handle advanced image augmentation techniques.
• 🖼️ Albumentations are now applied during training in utils/datasets.py if augmentations are enabled.
• 🧩 Added version check for Albumentations in utils/general.py.

🎯 Purpose & Impact

• 👁️‍🗨️ This PR introduces more complex and diverse image transformations during model training for improved robustness and potential accuracy enhancements.
• 🧠 The integration allows users to leverage a broader range of augmentation options, potentially leading to better model generalization and performance.
• 🔄 The change is optional and backwards compatible, ensuring current workflows are uninterrupted for users who do not wish to use this new feature.

[glenn-jocher]

/rebase

[glenn-jocher]

/rebase

[JosepBC]

Something like this if I just want mosaic, right?
hsv_h: 0.0 # image HSV-Hue augmentation (fraction)
hsv_s: 0.0 # image HSV-Saturation augmentation (fraction)
hsv_v: 0.0 # image HSV-Value augmentation (fraction)
degrees: 0.0 # image rotation (+/- deg)
translate: 0.0 # image translation (+/- fraction)
scale: 0.0 # image scale (+/- gain)
shear: 0.0 # image shear (+/- deg)
perspective: 0.0 # image perspective (+/- fraction), range 0-0.001
flipud: 0.0 # image flip up-down (probability)
fliplr: 0.0 # image flip left-right (probability)
mosaic: 1.0 # image mosaic (probability)
mixup: 0.0 # image mixup (probability)
copy_paste: 0.0 # segment copy-paste (probability)

[glenn-jocher]

yes