[Feature] Add support for mask diagonal flip in TTA (open-mmlab#5403)

* Add support for mask diagonal flip in tta

* Add unit test

* Fix unit test

* Fix unit test

Author: yeliudev

mmdet/core/post_processing/merge_augs.py

@@ -137,6 +137,9 @@ def merge_aug_masks(aug_masks, img_metas, rcnn_test_cfg, weights=None):
                 mask = mask[:, :, :, ::-1]
             elif flip_direction == 'vertical':
                 mask = mask[:, :, ::-1, :]
+            elif flip_direction == 'diagonal':
+                mask = mask[:, :, :, ::-1]
+                mask = mask[:, :, ::-1, :]
             else:
                 raise ValueError(
                     f"Invalid flipping direction '{flip_direction}'")

mmdet/datasets/pipelines/test_time_aug.py

@@ -44,9 +44,10 @@ class MultiScaleFlipAug:
         scale_factor (float | list[float] | None): Scale factors for resizing.
         flip (bool): Whether apply flip augmentation. Default: False.
         flip_direction (str | list[str]): Flip augmentation directions,
-            options are "horizontal" and "vertical". If flip_direction is list,
-            multiple flip augmentations will be applied.
-            It has no effect when flip == False. Default: "horizontal".
+            options are "horizontal", "vertical" and "diagonal". If
+            flip_direction is a list, multiple flip augmentations will be
+            applied. It has no effect when flip == False. Default:
+            "horizontal".
     """
 
     def __init__(self,

tests/test_data/test_pipelines/test_transform/test_models_aug_test.py

@@ -20,6 +20,7 @@ def model_aug_test_template(cfg_file):
     # init test pipeline and set aug test
     load_cfg, multi_scale_cfg = cfg.test_pipeline
     multi_scale_cfg['flip'] = True
+    multi_scale_cfg['flip_direction'] = ['horizontal', 'vertical', 'diagonal']
     multi_scale_cfg['img_scale'] = [(1333, 800), (800, 600), (640, 480)]
 
     load = build_from_cfg(load_cfg, PIPELINES)
@@ -29,8 +30,8 @@ def model_aug_test_template(cfg_file):
         img_prefix=osp.join(osp.dirname(__file__), '../../../data'),
         img_info=dict(filename='color.jpg'))
     results = transform(load(results))
-    assert len(results['img']) == 6
-    assert len(results['img_metas']) == 6
+    assert len(results['img']) == 12
+    assert len(results['img_metas']) == 12
 
     results['img'] = [collate([x]) for x in results['img']]
     results['img_metas'] = [collate([x]).data[0] for x in results['img_metas']]
@@ -56,14 +57,14 @@ def test_aug_test_size():
         transforms=[],
         img_scale=[(1333, 800), (800, 600), (640, 480)],
         flip=True,
-        flip_direction=['horizontal', 'vertical'])
+        flip_direction=['horizontal', 'vertical', 'diagonal'])
     multi_aug_test_module = build_from_cfg(transform, PIPELINES)
 
     results = load(results)
     results = multi_aug_test_module(load(results))
-    # len(["original", "horizontal", "vertical"]) *
+    # len(["original", "horizontal", "vertical", "diagonal"]) *
     # len([(1333, 800), (800, 600), (640, 480)])
-    assert len(results['img']) == 9
+    assert len(results['img']) == 12
 
 
 def test_cascade_rcnn_aug_test():
@@ -107,6 +108,7 @@ def test_cornernet_aug_test():
     # init test pipeline and set aug test
     load_cfg, multi_scale_cfg = cfg.test_pipeline
     multi_scale_cfg['flip'] = True
+    multi_scale_cfg['flip_direction'] = ['horizontal', 'vertical', 'diagonal']
     multi_scale_cfg['scale_factor'] = [0.5, 1.0, 2.0]
 
     load = build_from_cfg(load_cfg, PIPELINES)
@@ -116,8 +118,8 @@ def test_cornernet_aug_test():
         img_prefix=osp.join(osp.dirname(__file__), '../../../data'),
         img_info=dict(filename='color.jpg'))
     results = transform(load(results))
-    assert len(results['img']) == 6
-    assert len(results['img_metas']) == 6
+    assert len(results['img']) == 12
+    assert len(results['img_metas']) == 12
 
     results['img'] = [collate([x]) for x in results['img']]
     results['img_metas'] = [collate([x]).data[0] for x in results['img_metas']]