⚗️ [Add] MixUp augment, not sure it can work with Mosaic

config/data/augmentation.yaml

@@ -1,2 +1,3 @@
-RandomHorizontalFlip: 0.5
-Mosaic: 0.5
+Mosaic: 1
+MixUp: 1
+RandomHorizontalFlip: 0.5

utils/data_augment.py

@@ -2,6 +2,7 @@ from PIL import Image
 import numpy as np
 import torch
 from torchvision.transforms import functional as TF
+from torchvision.transforms.functional import to_tensor, to_pil_image
 
 
 class Compose:
@@ -77,3 +78,37 @@ class Mosaic:
 
         all_labels = torch.cat(all_labels, dim=0)
         return mosaic_image, all_labels
+
+
+class MixUp:
+    """Applies the MixUp augmentation to a pair of images and their corresponding boxes."""
+
+    def __init__(self, prob=0.5, alpha=1.0):
+        self.alpha = alpha
+        self.prob = prob
+        self.parent = None
+
+    def set_parent(self, parent):
+        """Set the parent dataset object for accessing dataset methods."""
+        self.parent = parent
+
+    def __call__(self, image, boxes):
+        if torch.rand(1) >= self.prob:
+            return image, boxes
+
+        assert self.parent is not None, "Parent is not set. MixUp cannot retrieve additional data."
+
+        # Retrieve another image and its boxes randomly from the dataset
+        image2, boxes2 = self.parent.get_more_data()[0]
+
+        # Calculate the mixup lambda parameter
+        lam = np.random.beta(self.alpha, self.alpha) if self.alpha > 0 else 0.5
+
+        # Mix images
+        image1, image2 = to_tensor(image), to_tensor(image2)
+        mixed_image = lam * image1 + (1 - lam) * image2
+
+        # Mix bounding boxes
+        mixed_boxes = torch.cat([lam * boxes, (1 - lam) * boxes2])
+
+        return to_pil_image(mixed_image), mixed_boxes

utils/dataloader.py

@@ -10,7 +10,7 @@ from tqdm.rich import tqdm
 import diskcache as dc
 from typing import Union
 from drawer import draw_bboxes
-from data_augment import Compose, RandomHorizontalFlip, Mosaic
+from data_augment import Compose, RandomHorizontalFlip, Mosaic, MixUp
 
 
 class YoloDataset(Dataset):