diff --git a/maskrcnn_benchmark/config/defaults.py b/maskrcnn_benchmark/config/defaults.py index 65fbdaddd..ea5d4df73 100644 --- a/maskrcnn_benchmark/config/defaults.py +++ b/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 diff --git a/maskrcnn_benchmark/data/transforms/build.py b/maskrcnn_benchmark/data/transforms/build.py index 52385ea7d..e559c891e 100644 --- a/maskrcnn_benchmark/data/transforms/build.py +++ b/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, ] diff --git a/maskrcnn_benchmark/data/transforms/transforms.py b/maskrcnn_benchmark/data/transforms/transforms.py index 2d37dc72f..ac8039c75 100644 --- a/maskrcnn_benchmark/data/transforms/transforms.py +++ b/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, diff --git a/maskrcnn_benchmark/structures/bounding_box.py b/maskrcnn_benchmark/structures/bounding_box.py index 4084024fa..96f7ec227 100644 --- a/maskrcnn_benchmark/structures/bounding_box.py +++ b/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