🔀 [Merge] branch 'DATASET' of github.com:LucyTuan/yolov9mit into DATASET

requirements.txt

@@ -6,4 +6,7 @@ pyyaml
 requests
 rich
 torch
-tqdm
+tqdm
+Pillow
+diskcache
+torchvision

utils/data_augment.py

@@ -37,6 +37,19 @@ class RandomHorizontalFlip:
         return image, boxes
 
 
+class RandomVerticalFlip:
+    """Randomly vertically flips the image along with the bounding boxes."""
+
+    def __init__(self, prob=0.5):
+        self.prob = prob
+
+    def __call__(self, image, boxes):
+        if torch.rand(1) < self.prob:
+            image = TF.vflip(image)
+            boxes[:, [2, 4]] = 1 - boxes[:, [4, 2]]
+        return image, boxes
+
+
 class Mosaic:
     """Applies the Mosaic augmentation to a batch of images and their corresponding 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, MixUp
+from data_augment import Compose, RandomHorizontalFlip, RandomVerticalFlip, Mosaic, MixUp
 
 
 class YoloDataset(Dataset):