Feature Request: Albumentations for YOLOv5 🚀 TTA

[glenn-jocher]

🚀 Feature

YOLOv5 uses 3 fixed transforms during Test Time Augmentations (TTA): 3 scales with a left-right flip on the middle scale. Our users are requesting specific transforms for TTA (i.e. rot90) and proposing Qubvel integration for this, but I think it would be more useful to extend the Albumentations capability to TTA to allow for any combination of transforms. It seems like the major capability that's lacking is inverting the predictions (using the inverse transform).

Motivation

Motivated by multiple YOLOv5 user requests:

ultralytics/yolov5#4107
ultralytics/yolov5#4394

Pitch

A possible high-level pseudocode workflow might be:

for xform in tta_transforms:
    im_x = xform(im, seed=0)  # Albumentations transform
    results_x = model(im_x)  # inference
    results.append(xform.invert_labels(results_x, seed=0))  # <--- invert predictions

predictions = NMS(torch.stack(results))  # fused predictions from an arbitrary set of albumentations

Alternatives

Qubvel integration with YOLOv5 or extending the YOLOv5 TTA code itself to support more transforms.

Additional context

This would be a natural followon to our excellent Albumentations integration for training augmentation: #949

@BloodAxe

[Dipet]

Hi! I started some work on TTA, you can see how it might look in this PR: #974
I will try to speed up the implementation of this function.

[RachelRamirez]

I have also posted an issue/feature request/example request to improve TTA here: albumentations-team/albumentations_examples#25