color jitter only during training (facebookresearch#703)

Lyears · Apr 20, 2019 · fe68c41 · fe68c41
1 parent c9210b1
commit fe68c41
Showing 1 changed file with 13 additions and 5 deletions.
diff --git a/maskrcnn_benchmark/data/transforms/build.py b/maskrcnn_benchmark/data/transforms/build.py
@@ -7,21 +7,29 @@ def build_transforms(cfg, is_train=True):
         min_size = cfg.INPUT.MIN_SIZE_TRAIN
         max_size = cfg.INPUT.MAX_SIZE_TRAIN
         flip_prob = 0.5  # cfg.INPUT.FLIP_PROB_TRAIN
+        brightness = cfg.INPUT.BRIGHTNESS
+        contrast = cfg.INPUT.CONTRAST
+        saturation = cfg.INPUT.SATURATION
+        hue = cfg.INPUT.HUE
     else:
         min_size = cfg.INPUT.MIN_SIZE_TEST
         max_size = cfg.INPUT.MAX_SIZE_TEST
         flip_prob = 0
+        brightness = 0.0
+        contrast = 0.0
+        saturation = 0.0
+        hue = 0.0
 
     to_bgr255 = cfg.INPUT.TO_BGR255
     normalize_transform = T.Normalize(
         mean=cfg.INPUT.PIXEL_MEAN, std=cfg.INPUT.PIXEL_STD, to_bgr255=to_bgr255
     )
     color_jitter = T.ColorJitter(
-        brightness=cfg.INPUT.BRIGHTNESS,
-        contrast=cfg.INPUT.CONTRAST,
-        saturation=cfg.INPUT.SATURATION,
-        hue=cfg.INPUT.HUE,
-        )
+        brightness=brightness,
+        contrast=contrast,
+        saturation=saturation,
+        hue=hue,
+    )
 
     transform = T.Compose(
         [