yjh0410 2 年之前
父节点
当前提交
b3bc52b6e5
共有 6 个文件被更改,包括 56 次插入68 次删除
  1. 1 1
      config/data_config/transform_config.py
  2. 2 0
      dataset/build.py
  3. 16 15
      dataset/coco.py
  4. 10 9
      dataset/data_augment/yolov5_augment.py
  5. 10 10
      dataset/ourdataset.py
  6. 17 33
      dataset/voc.py

+ 1 - 1
config/data_config/transform_config.py

@@ -199,7 +199,7 @@ yolox_nano_trans_config = {
     # Basic Augment
     'degrees': 10.0,
     'translate': 0.1,
-    'scale': 0.9,
+    'scale': 0.5,
     'shear': 2.0,
     'perspective': 0.0,
     'hsv_h': 0.015,

+ 2 - 0
dataset/build.py

@@ -81,6 +81,8 @@ def build_transform(args, trans_config, max_stride=32, is_train=False):
             transform = SSDAugmentation(img_size=args.img_size,)
         else:
             transform = SSDBaseTransform(img_size=args.img_size,)
+        trans_config['mosaic_prob'] = 0.0
+        trans_config['mixup_prob'] = 0.0
 
     elif trans_config['aug_type'] == 'yolov5':
         if is_train:

+ 16 - 15
dataset/coco.py

@@ -226,15 +226,21 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description='COCO-Dataset')
 
     # opt
-    parser.add_argument('--root', default='D:\\python_work\\object-detection\\dataset\\COCO',
+    parser.add_argument('--root', default='/Users/liuhaoran/Desktop/python_work/object-detection/dataset/COCO/',
                         help='data root')
+    parser.add_argument('-size', '--img_size', default=640, type=int,
+                        help='input image size.')
+    parser.add_argument('--mosaic', default=None, type=float,
+                        help='mosaic augmentation.')
+    parser.add_argument('--mixup', default=None, type=float,
+                        help='mixup augmentation.')
+    parser.add_argument('--is_train', action="store_true", default=False,
+                        help='mixup augmentation.')
     
     args = parser.parse_args()
 
-    is_train = False
-    img_size = 640
-    yolov5_trans_config = {
-        'aug_type': 'yolov5',
+    trans_config = {
+        'aug_type': 'yolov5',  # optional: ssd, yolov5
         # Basic Augment
         'degrees': 0.0,
         'translate': 0.2,
@@ -246,26 +252,21 @@ if __name__ == "__main__":
         'hsv_v': 0.4,
         # Mosaic & Mixup
         'mosaic_prob': 1.0,
-        'mixup_prob': 0.15,
+        'mixup_prob': 1.0,
         'mosaic_type': 'yolov5_mosaic',
         'mixup_type': 'yolov5_mixup',
         'mixup_scale': [0.5, 1.5]
     }
-    ssd_trans_config = {
-        'aug_type': 'ssd',
-        'mosaic_prob': 0.0,
-        'mixup_prob': 0.0
-    }
 
-    transform = build_transform(img_size, yolov5_trans_config, is_train)
+    transform, trans_cfg = build_transform(args, trans_config, 32, args.is_train)
 
     dataset = COCODataset(
-        img_size=img_size,
+        img_size=args.img_size,
         data_dir=args.root,
         image_set='val2017',
-        trans_config=yolov5_trans_config,
+        trans_config=trans_config,
         transform=transform,
-        is_train=is_train
+        is_train=args.is_train
         )
     
     np.random.seed(0)

+ 10 - 9
dataset/data_augment/yolov5_augment.py

@@ -5,6 +5,8 @@ import numpy as np
 import torch
 
 
+# ------------------------- Basic augmentations -------------------------
+## Spatial transform
 def random_perspective(image,
                        targets=(),
                        degrees=10,
@@ -78,7 +80,7 @@ def random_perspective(image,
 
     return image, targets
 
-
+## Color transform
 def augment_hsv(img, hgain=0.5, sgain=0.5, vgain=0.5):
     r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1  # random gains
     hue, sat, val = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
@@ -93,7 +95,8 @@ def augment_hsv(img, hgain=0.5, sgain=0.5, vgain=0.5):
     cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR, dst=img)  # no return needed
 
 
-# YOLOv5-Mosaic
+# ------------------------- Strong augmentations -------------------------
+## YOLOv5-Mosaic
 def yolov5_mosaic_augment(image_list, target_list, img_size, affine_params=None, is_train=False):
     assert len(image_list) == 4
 
@@ -180,8 +183,7 @@ def yolov5_mosaic_augment(image_list, target_list, img_size, affine_params=None,
 
     return mosaic_img, mosaic_target
 
-
-# YOLOv5-Mixup
+## YOLOv5-Mixup
 def yolov5_mixup_augment(origin_image, origin_target, new_image, new_target):
     if origin_image.shape[:2] != new_image.shape[:2]:
         img_size = max(new_image.shape[:2])
@@ -218,8 +220,7 @@ def yolov5_mixup_augment(origin_image, origin_target, new_image, new_target):
     
     return mixup_image, mixup_target
     
-
-# YOLOX-Mixup
+## YOLOX-Mixup
 def yolox_mixup_augment(origin_img, origin_target, new_img, new_target, img_size, mixup_scale):
     jit_factor = random.uniform(*mixup_scale)
     FLIP = random.uniform(0, 1) > 0.5
@@ -299,7 +300,8 @@ def yolox_mixup_augment(origin_img, origin_target, new_img, new_target, img_size
     return origin_img.astype(np.uint8), mixup_target
         
 
-# YOLOv5-style TrainTransform
+# ------------------------- Preprocessers -------------------------
+## YOLOv5-style Transform for Train
 class YOLOv5Augmentation(object):
     def __init__(self, 
                  img_size=640,
@@ -374,8 +376,7 @@ class YOLOv5Augmentation(object):
 
         return pad_image, target, [dw, dh]
 
-
-# YOLOv5-style ValTransform
+## YOLOv5-style Transform for Eval
 class YOLOv5BaseTransform(object):
     def __init__(self, img_size=640, max_stride=32):
         self.img_size = img_size

+ 10 - 10
dataset/ourdataset.py

@@ -203,7 +203,7 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description='FreeYOLOv2')
 
     # opt
-    parser.add_argument('--root', default='AnimalDataset',
+    parser.add_argument('--root', default='/Users/liuhaoran/Desktop/python_work/object-detection/dataset/AnimalDataset/',
                         help='data root')
     parser.add_argument('--split', default='train',
                         help='data split')
@@ -215,13 +215,13 @@ if __name__ == "__main__":
                         help='mosaic augmentation.')
     parser.add_argument('--mixup', default=None, type=float,
                         help='mixup augmentation.')
-
-    args = parser.parse_args()
+    parser.add_argument('--is_train', action="store_true", default=False,
+                        help='mixup augmentation.')
     
-    img_size = 640
-    is_train = True
+    args = parser.parse_args()
+
     trans_config = {
-        'aug_type': 'yolov5',
+        'aug_type': 'yolov5',  # optional: ssd, yolov5
         # Basic Augment
         'degrees': 0.0,
         'translate': 0.2,
@@ -233,16 +233,16 @@ if __name__ == "__main__":
         'hsv_v': 0.4,
         # Mosaic & Mixup
         'mosaic_prob': 1.0,
-        'mosaic_9x_prob': 0.2,
-        'mixup_prob': 0.15,
+        'mixup_prob': 1.0,
         'mosaic_type': 'yolov5_mosaic',
         'mixup_type': 'yolov5_mixup',
         'mixup_scale': [0.5, 1.5]
     }
-    transform, trans_config = build_transform(args, trans_config, max_stride=32, is_train=is_train)
+
+    transform, trans_cfg = build_transform(args, trans_config, 32, args.is_train)
 
     dataset = OurDataset(
-        img_size=img_size,
+        img_size=args.img_size,
         data_dir=args.root,
         image_set=args.split,
         transform=transform,

+ 17 - 33
dataset/voc.py

@@ -249,33 +249,22 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description='VOC-Dataset')
 
     # opt
-    parser.add_argument('--root', default='D:\\python_work\\object-detection\\dataset\\VOCdevkit',
+    parser.add_argument('--root', default='/Users/liuhaoran/Desktop/python_work/object-detection/dataset/VOCdevkit/',
                         help='data root')
+    parser.add_argument('-size', '--img_size', default=640, type=int,
+                        help='input image size.')
+    parser.add_argument('--mosaic', default=None, type=float,
+                        help='mosaic augmentation.')
+    parser.add_argument('--mixup', default=None, type=float,
+                        help='mixup augmentation.')
+    parser.add_argument('--is_train', action="store_true", default=False,
+                        help='mixup augmentation.')
+    
     
     args = parser.parse_args()
 
-    is_train = False
-    img_size = 640
-    yolov5_trans_config = {
-        'aug_type': 'yolov5',
-        # Basic Augment
-        'degrees': 0.0,
-        'translate': 0.2,
-        'scale': 0.9,
-        'shear': 0.0,
-        'perspective': 0.0,
-        'hsv_h': 0.015,
-        'hsv_s': 0.7,
-        'hsv_v': 0.4,
-        # Mosaic & Mixup
-        'mosaic_prob': 1.0,
-        'mixup_prob': 0.15,
-        'mosaic_type': 'yolov5_mosaic',
-        'mixup_type': 'yolov5_mixup',
-        'mixup_scale': [0.5, 1.5]
-    }
-    yolox_trans_config = {
-        'aug_type': 'yolov5',
+    trans_config = {
+        'aug_type': 'yolov5',  # optional: ssd, yolov5
         # Basic Augment
         'degrees': 0.0,
         'translate': 0.2,
@@ -289,22 +278,17 @@ if __name__ == "__main__":
         'mosaic_prob': 1.0,
         'mixup_prob': 1.0,
         'mosaic_type': 'yolov5_mosaic',
-        'mixup_type': 'yolox_mixup',
+        'mixup_type': 'yolov5_mixup',
         'mixup_scale': [0.5, 1.5]
     }
-    ssd_trans_config = {
-        'aug_type': 'ssd',
-        'mosaic_prob': 0.0,
-        'mixup_prob': 0.0
-    }
-    transform = build_transform(img_size, yolov5_trans_config, is_train)
+    transform, trans_cfg = build_transform(args, trans_config, 32, args.is_train)
 
     dataset = VOCDetection(
-        img_size=img_size,
+        img_size=args.img_size,
         data_dir=args.root,
-        trans_config=yolov5_trans_config,
+        trans_config=trans_config,
         transform=transform,
-        is_train=is_train
+        is_train=args.is_train
         )
     
     np.random.seed(0)