henry000 commited on
Commit
5ce12fa
Β·
2 Parent(s): 49d58b9 120ff32

πŸ”€ [Merge] branch 'DATASET' of github.com:LucyTuan/yolov9mit into DATASET

Browse files
Files changed (3) hide show
  1. requirements.txt +4 -1
  2. utils/data_augment.py +13 -0
  3. utils/dataloader.py +1 -1
requirements.txt CHANGED
@@ -6,4 +6,7 @@ pyyaml
6
  requests
7
  rich
8
  torch
9
- tqdm
 
 
 
 
6
  requests
7
  rich
8
  torch
9
+ tqdm
10
+ Pillow
11
+ diskcache
12
+ torchvision
utils/data_augment.py CHANGED
@@ -37,6 +37,19 @@ class RandomHorizontalFlip:
37
  return image, boxes
38
 
39
 
 
 
 
 
 
 
 
 
 
 
 
 
 
40
  class Mosaic:
41
  """Applies the Mosaic augmentation to a batch of images and their corresponding boxes."""
42
 
 
37
  return image, boxes
38
 
39
 
40
+ class RandomVerticalFlip:
41
+ """Randomly vertically flips the image along with the bounding boxes."""
42
+
43
+ def __init__(self, prob=0.5):
44
+ self.prob = prob
45
+
46
+ def __call__(self, image, boxes):
47
+ if torch.rand(1) < self.prob:
48
+ image = TF.vflip(image)
49
+ boxes[:, [2, 4]] = 1 - boxes[:, [4, 2]]
50
+ return image, boxes
51
+
52
+
53
  class Mosaic:
54
  """Applies the Mosaic augmentation to a batch of images and their corresponding boxes."""
55
 
utils/dataloader.py CHANGED
@@ -10,7 +10,7 @@ from tqdm.rich import tqdm
10
  import diskcache as dc
11
  from typing import Union
12
  from drawer import draw_bboxes
13
- from data_augment import Compose, RandomHorizontalFlip, Mosaic, MixUp
14
 
15
 
16
  class YoloDataset(Dataset):
 
10
  import diskcache as dc
11
  from typing import Union
12
  from drawer import draw_bboxes
13
+ from data_augment import Compose, RandomHorizontalFlip, RandomVerticalFlip, Mosaic, MixUp
14
 
15
 
16
  class YoloDataset(Dataset):