add probability flag for hsv_augmentation and random_flip (#592)

exps/example/yolox_voc/yolox_voc_s.py

@@ -15,6 +15,13 @@ class Exp(MyExp):
         self.depth = 0.33
         self.width = 0.50
         self.warmup_epochs = 1
+
+        # ---------- transform config ------------ #
+        self.mosaic_prob = 1.0
+        self.mixup_prob = 1.0
+        self.hsv_prob = 1.0
+        self.flip_prob = 0.5
+
         self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
 
     def get_data_loader(self, batch_size, is_distributed, no_aug=False, cache_img=False):
@@ -38,7 +45,10 @@ class Exp(MyExp):
                 data_dir=os.path.join(get_yolox_datadir(), "VOCdevkit"),
                 image_sets=[('2007', 'trainval'), ('2012', 'trainval')],
                 img_size=self.input_size,
-                preproc=TrainTransform(max_labels=50),
+                preproc=TrainTransform(
+                    max_labels=50,
+                    flip_prob=self.flip_prob,
+                    hsv_prob=self.hsv_prob),
                 cache=cache_img,
             )
 
@@ -46,7 +56,10 @@ class Exp(MyExp):
             dataset,
             mosaic=not no_aug,
             img_size=self.input_size,
-            preproc=TrainTransform(max_labels=120),
+            preproc=TrainTransform(
+                max_labels=120,
+                flip_prob=self.flip_prob,
+                hsv_prob=self.hsv_prob),
             degrees=self.degrees,
             translate=self.translate,
             mosaic_scale=self.mosaic_scale,

yolox/data/data_augment.py

@@ -140,9 +140,9 @@ def random_perspective(
     return img, targets
 
 
-def _mirror(image, boxes):
+def _mirror(image, boxes, prob=0.5):
     _, width, _ = image.shape
-    if random.randrange(2):
+    if random.random() < prob:
         image = image[:, ::-1]
         boxes = boxes.copy()
         boxes[:, 0::2] = width - boxes[:, 2::-2]
@@ -169,8 +169,10 @@ def preproc(img, input_size, swap=(2, 0, 1)):
 
 
 class TrainTransform:
-    def __init__(self, max_labels=50):
+    def __init__(self, max_labels=50, flip_prob=0.5, hsv_prob=1.0):
         self.max_labels = max_labels
+        self.flip_prob = flip_prob
+        self.hsv_prob = hsv_prob
 
     def __call__(self, image, targets, input_dim):
         boxes = targets[:, :4].copy()
@@ -188,8 +190,9 @@ class TrainTransform:
         # bbox_o: [xyxy] to [c_x,c_y,w,h]
         boxes_o = xyxy2cxcywh(boxes_o)
 
-        augment_hsv(image)
-        image_t, boxes = _mirror(image, boxes)
+        if random.random() < self.hsv_prob:
+            augment_hsv(image)
+        image_t, boxes = _mirror(image, boxes, self.flip_prob)
         height, width, _ = image_t.shape
         image_t, r_ = preproc(image_t, input_dim)
         # boxes [xyxy] 2 [cx,cy,w,h]

yolox/exp/yolox_base.py

@@ -38,6 +38,8 @@ class Exp(BaseExp):
         # --------------- transform config ----------------- #
         self.mosaic_prob = 1.0
         self.mixup_prob = 1.0
+        self.hsv_prob = 1.0
+        self.flip_prob = 0.5
         self.degrees = 10.0
         self.translate = 0.1
         self.mosaic_scale = (0.1, 2)
@@ -110,7 +112,10 @@ class Exp(BaseExp):
                 data_dir=self.data_dir,
                 json_file=self.train_ann,
                 img_size=self.input_size,
-                preproc=TrainTransform(max_labels=50),
+                preproc=TrainTransform(
+                    max_labels=50,
+                    flip_prob=self.flip_prob,
+                    hsv_prob=self.hsv_prob),
                 cache=cache_img,
             )
 
@@ -118,7 +123,10 @@ class Exp(BaseExp):
             dataset,
             mosaic=not no_aug,
             img_size=self.input_size,
-            preproc=TrainTransform(max_labels=120),
+            preproc=TrainTransform(
+                max_labels=120,
+                flip_prob=self.flip_prob,
+                hsv_prob=self.hsv_prob),
             degrees=self.degrees,
             translate=self.translate,
             mosaic_scale=self.mosaic_scale,