add rotation augmentation

maskrcnn_benchmark/config/defaults.py

@@ -61,6 +61,7 @@
 _C.INPUT.HUE = 0.0
 
 _C.INPUT.VERTICAL_FLIP_PROB_TRAIN = 0.0
+_C.INPUT.ROTATE90_PROB_TRAIN = 0.0
 
 # -----------------------------------------------------------------------------
 # Dataset

maskrcnn_benchmark/data/transforms/build.py

@@ -8,6 +8,7 @@ def build_transforms(cfg, is_train=True):
         max_size = cfg.INPUT.MAX_SIZE_TRAIN
         flip_horizontal_prob = 0.5  # cfg.INPUT.FLIP_PROB_TRAIN
         flip_vertical_prob = cfg.INPUT.VERTICAL_FLIP_PROB_TRAIN
+        rot_prob = cfg.INPUT.ROTATE90_PROB_TRAIN
         brightness = cfg.INPUT.BRIGHTNESS
         contrast = cfg.INPUT.CONTRAST
         saturation = cfg.INPUT.SATURATION
@@ -17,6 +18,7 @@ def build_transforms(cfg, is_train=True):
         max_size = cfg.INPUT.MAX_SIZE_TEST
         flip_horizontal_prob = 0.0
         flip_vertical_prob = 0.0
+        rot_prob = 0.0
         brightness = 0.0
         contrast = 0.0
         saturation = 0.0
@@ -39,6 +41,7 @@ def build_transforms(cfg, is_train=True):
             T.Resize(min_size, max_size),
             T.RandomHorizontalFlip(flip_horizontal_prob),
             T.RandomVerticalFlip(flip_vertical_prob),
+            T.RandomRotate90(rot_prob),
             T.ToTensor(),
             normalize_transform,
         ]

maskrcnn_benchmark/data/transforms/transforms.py

@@ -82,6 +82,18 @@ def __call__(self, image, target):
             image = F.vflip(image)
             target = target.transpose(1)
         return image, target
+class RandomRotate90(object):
+    def __init__(self, prob=0.5):
+        self.prob = prob
+
+    def __call__(self, image, target):
+        if random.random() < self.prob:
+            oriSize = image.size
+            image = F.rotate(image,90,expand=True)
+            image = F.resize(image, oriSize)
+
+            target = target.transpose(2)
+        return image, target
 
 class ColorJitter(object):
     def __init__(self,

maskrcnn_benchmark/structures/bounding_box.py

@@ -4,7 +4,7 @@
 # transpose
 FLIP_LEFT_RIGHT = 0
 FLIP_TOP_BOTTOM = 1
-
+ROTATE_90 = 2
 
 class BoxList(object):
     """
@@ -152,6 +152,11 @@ def transpose(self, method):
             transposed_xmax = xmax
             transposed_ymin = image_height - ymax
             transposed_ymax = image_height - ymin
+        elif method == ROTATE_90:
+            transposed_xmin = ymin * image_width/image_height
+            transposed_xmax = ymax * image_width/image_height
+            transposed_ymin = (image_width - xmax) * image_height / image_width
+            transposed_ymax = (image_width - xmin) * image_height/ image_width
 
         transposed_boxes = torch.cat(
             (transposed_xmin, transposed_ymin, transposed_xmax, transposed_ymax), dim=-1