How to incorporate custom data augmentation during training?

[BoPengGit]

❔Question

Hi Glenn,

Great work. I know in Kaggle for the wheat competition many months ago YoloV5 was a great hit.

I'm wanting to use custom augmentation during training for my dataset. How can I do that? Is this the training augmentation config file?

I'm wanting to add: RandomContrast, RandomGamma, RandomBrightness, ElasticTransform, GridDistortion, OpticalDistortion, ShiftScaleRotate.

Thanks Glenn,

Here's my Kaggle Profile: www.kaggle.com/bopengiowa

[github-actions]

👋 Hello @BoPengGit, thank you for your interest in 🚀 YOLOv5! Please visit our ⭐️ Tutorials to get started, where you can find quickstart guides for simple tasks like Custom Data Training all the way to advanced concepts like Hyperparameter Evolution.

If this is a 🐛 Bug Report, please provide screenshots and minimum viable code to reproduce your issue, otherwise we can not help you.

If this is a custom training ❓ Question, please provide as much information as possible, including dataset images, training logs, screenshots, and a public link to online W&B logging if available.

For business inquiries or professional support requests please visit https://www.ultralytics.com or email Glenn Jocher at glenn.jocher@ultralytics.com.

Requirements

Python 3.8 or later with all requirements.txt dependencies installed, including torch>=1.7. To install run:

$ pip install -r requirements.txt

Environments

YOLOv5 may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

• Google Colab Notebook with free GPU:
• Kaggle Notebook with free GPU: https://www.kaggle.com/models/ultralytics/yolov5
• Google Cloud Deep Learning VM. See GCP Quickstart Guide
• Docker Image https://hub.docker.com/r/ultralytics/yolov5. See Docker Quickstart Guide

Status

If this badge is green, all YOLOv5 GitHub Actions Continuous Integration (CI) tests are currently passing. CI tests verify correct operation of YOLOv5 training (train.py), testing (test.py), inference (detect.py) and export (export.py) on MacOS, Windows, and Ubuntu every 24 hours and on every commit.

[glenn-jocher]

@BoPengGit yes that's the correct location. The image augmentation hyperparameters are here:

yolov5/data/hyp.scratch.yaml

Lines 22 to 33 in b75c432

	hsv_h: 0.015 # image HSV-Hue augmentation (fraction)
	hsv_s: 0.7 # image HSV-Saturation augmentation (fraction)
	hsv_v: 0.4 # image HSV-Value augmentation (fraction)
	degrees: 0.0 # image rotation (+/- deg)
	translate: 0.1 # image translation (+/- fraction)
	scale: 0.5 # 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.5 # image flip left-right (probability)
	mosaic: 1.0 # image mosaic (probability)
	mixup: 0.0 # image mixup (probability)

If you wanted to add additional hyperparameters you could add them in the hyp file and then access them here in the dataloader (i.e. to place an albumentations call to the more exotic augmentations). This would make an interesting PR btw, others have also expressed interest in additional augmentations.

yolov5/utils/datasets.py

Lines 495 to 508 in b75c432

	hyp = self.hyp
	mosaic = self.mosaic and random.random() < hyp['mosaic']
	if mosaic:
	# Load mosaic
	img, labels = load_mosaic(self, index)
	shapes = None
	# MixUp https://arxiv.org/pdf/1710.09412.pdf
	if random.random() < hyp['mixup']:
	img2, labels2 = load_mosaic(self, random.randint(0, self.n - 1))
	r = np.random.beta(8.0, 8.0) # mixup ratio, alpha=beta=8.0
	img = (img * r + img2 * (1 - r)).astype(np.uint8)
	labels = np.concatenate((labels, labels2), 0)

[glenn-jocher]

@BoPengGit one item I should mention is that the image-space augmentations (default HSV augmentations and the RandomContrast, RandomGamma, RandomBrightness you mentioned) are currently implemented a single time for an entire mosaic. It would be interesting to see if implementing them per each image in the mosaic helps or hurts. The HSV augmentations are implemented in same region of the dataloader here:

yolov5/utils/datasets.py

Lines 539 to 540 in b75c432

	# Augment colorspace
	augment_hsv(img, hgain=hyp['hsv_h'], sgain=hyp['hsv_s'], vgain=hyp['hsv_v'])

[BoPengGit]

Hi Gleen,

I think I understand what your saying. The HSV augmentation is after the mosaic augmentation in the code. Is this what your saying?

Also for the augment_hsv() method listed above, that method doesn't return anything correct? And inside the method implementation there is no self. instance variables that are modified. So does that method have any affect on the image?

I understand, so in that code space you want me to write code for augmentation like ElasticTransform, GridDistortion, OpticalDistortion.

Can you mention which of the RandomContrast, RandomGamma, RandomBrightness, ElasticTransform, GridDistortion, OpticalDistortion, ShiftScaleRotate augmentations are currently implemented?

Thanks Gleen,

[glenn-jocher]

@BoPengGit augment_hsv(img) will modify img inplace, so while the function does not return anything, img outside the function will be modified correctly after line 540.

The high level augmentation overview is here, you can see augment_hsv() working correctly, modifying an entire image (and background). I have not tested image-space (i.e. HSV) augmentation applied to individual images instead of entire mosaics. This may add more randomness so I would assume this would be a higher level of augmentation, though the background would never change, so this reduce background robustness a bit.

HSV augmentation may be redundant with the 3 you posted there, or they may be complementary, I'm not sure. Saturation introduces color, i.e. zero saturation produces a greyscale image I believe, and Value is essentially the brightness of the image. The best way to do all this is just to implement the modifications and then pretend to start a training, i.e. python train.py --data coco128.yaml --epochs 1, which will start a new training run and plot the first few batches of mosaics in your runs/train/exp directory. You allow a few batches to train then examine the train*.jpgs to see the effect of the changes you made, and then keep iterating and reviewing the train*.jpgs to verify your implementation.

[glenn-jocher]

These images will directly show the effect of your augmentation updates once you start a training:

[glenn-jocher]

I think of the transforms you mentioned ShiftScaleRotate is the primary one already implemented, along with RandomBrightness. ShiftScaleRotate is implemented as a perspective transform here, composed of translation, scale, shift, shear and perspective (an affine transform is translation, scale, shift, shear but no perspective):

yolov5/utils/datasets.py

Lines 828 to 829 in b75c432

	def random_perspective(img, targets=(), degrees=10, translate=.1, scale=.1, shear=10, perspective=0.0, border=(0, 0)):
	# torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(.1, .1), scale=(.9, 1.1), shear=(-10, 10))

[BoPengGit]

Hi Gleen,

That sounds good. Thanks for that.

So can we just import the albumentations library and do something like a **kwargs key:value input into an albumentation function to do bounding box augmentation using the albumentations library?

The keys of the **kwargs input would be the augmentation names and the values of the **kwargs input would be the corresponding augmentation parameter values.

Can it be as simple as that?

Thanks Gleen,

[glenn-jocher]

@BoPengGit yes I think it may actually be pretty simple. I've not used albumentations myself, so I'm not sure if the label and image formats are directly passable to the package or if they need modification. If you get this working and would like to contribute a PR for a basic albumentations interface, that might be useful for future users, as I know several have expressed interest or asked about albumentations specifically. We may want to profile changes though to ensure they don't add any significant slowdowns to the existing dataloader (which has itself undergone much profiling and optimization for speed).

[BoPengGit]

Sounds good Gleen. Below is a sample albumentations agumentation for object detection from this notebook: https://www.kaggle.com/ar2017/efficientdet-train-mixup-cutmix-stratifiedk-fold

I'll probably try to get it working and can make a PR for you to test if I can get it working. I'm competing in this Kaggle competition https://www.kaggle.com/c/vinbigdata-chest-xray-abnormalities-detection so I may wait until competition is over to make a PR if I get it working. ;)

Or obviously you can make a PR as well.

def get_train_transforms():
    return A.Compose(
        [
            A.RandomSizedCrop(min_max_height=(500, 720), height=720, width=720, p=0.5),
            A.OneOf([
                A.HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit= 0.2, 
                                     val_shift_limit=0.2, p=0.9),
                A.RandomBrightnessContrast(brightness_limit=0.2, 
                                           contrast_limit=0.2, p=0.9),
            ], p=0.9),
            A.Resize(height=512, width=512, p=1),
            A.Cutout(num_holes=8, max_h_size=64, max_w_size=64, fill_value=0, p=0.5),
            ToTensorV2(p=1.0),
        ], 
        p=1.0, 
        bbox_params=A.BboxParams(
            format='pascal_voc',
            min_area=0, 
            min_visibility=0,
            label_fields=['labels']
        )
    )
...
   self.transforms = get_train_transforms()
    sample = self.transforms(**{
        'image': image,
        'bboxes': target['boxes'],
        'labels': labels
    })

[glenn-jocher]

Understood, sounds good! You may want to train and eval at a higher resolution (i.e. 1280 or native resolution if higher) with perhaps a p6 model here for best competition results:
https://github.com/ultralytics/yolov5/blob/master/models/hub/yolov5-p6.yaml

[BoPengGit]

Sounds good Gleen. Thanks for the suggestion on the model size.

I'll keep you posted on any updates on the albumentations side.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

[glenn-jocher]

@BoPengGit see PR #3882 for a proposed automatic Albumentations integration.

[glenn-jocher]

@BoPengGit good news 😃! Your original issue may now be fixed ✅ in PR #3882. This PR implements a YOLOv5 🚀 + Albumentations integration. The integration will automatically apply Albumentations transforms during YOLOv5 training if albumentations>=1.0.0 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 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.

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.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

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 spotting this issue and informing us of the problem. Please let us know if this update resolves the issue for you, and feel free to inform us of any other issues you discover or feature requests that come to mind. Happy trainings with YOLOv5 🚀!

[zakharchik-me]

@glenn-jocher thank you for your replies, they were helpful! One more question, is it possible to control number of augmented images.
As you had 3x3 matrix of augmented img.
For ex: make the matrix 20x20
And add an amount of modified images instead of one, to increase the count of low classes in dataset.

[glenn-jocher]

@zakharchik-me glad to hear that the previous replies were helpful! Regarding your question, currently the number of augmented images is determined by the batch size and the augmentations applied. However, you can potentially increase the number of modified images by adjusting the probability of applying certain augmentations such as random cropping and flipping. Alternatively, if you would like a more precise control over the number of augmented images, you may consider using a separate data augmentation library or modifying the augmentation pipeline in the existing codebase. I hope this helps! Let me know if you have any further questions.