From 20fa12912fba77732c0da3ec34663c8522d7606b Mon Sep 17 00:00:00 2001
From: Hongyuan Zhang <66273343+Alias-z@users.noreply.github.com>
Date: Fri, 14 Jul 2023 06:44:03 +0200
Subject: [PATCH] [Fix] Albumentations default key mapping mismatch (#3195)

## Modification

Fix Albumentations default key mapping mismatch as mentioned in [issue #
3179](https://github.com/open-mmlab/mmsegmentation/issues/3179) by
changing `self.keymap_to_albu = { 'img': 'image', 'gt_masks': 'masks'}`
to `self.keymap_to_albu = { 'img': 'image', 'gt_seg_map': 'mask'}`

## Use cases (Optional)

Example albu config
```
crop_size = (512, 512)

albu_train_transforms = [
    dict(
        type='PadIfNeeded',
        min_height=crop_size[0]*2,
        min_width=crop_size[1]*2,
        border_mode=0,
        always_apply=True),
    dict(type='Flip', always_apply=True),
    dict(type='Rotate', limit=(-180, 180), interpolation=4, always_apply=True),
    dict(type='RandomScale', scale_limit=0.1, interpolation=4, always_apply=True),
    dict(
        type='ElasticTransform',
        alpha=20,
        sigma=15,
        interpolation=4,
        border_mode=0,
        mask_value=(0, 0, 0),
        approximate=True,
        same_dxdy=True,
        p=0.8),
    dict(type='ColorJitter', brightness=0.2, contrast=0.1, saturation=0.2, hue=0.2, always_apply=True),
    dict(type='AdvancedBlur', p=0.5),
    dict(type='CenterCrop', height=crop_size[0], width=crop_size[1], always_apply=True)
]
```

Example training pipeline without specifying `keymap`
```
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations', reduce_zero_label=False),
    dict(
        type='Albu',
        transforms=albu_train_transforms,
        ),
    dict(type='Resize', scale=crop_size, keep_ratio=False, interpolation='lanczos'),
    dict(type='PackSegInputs')
]
```

Example viz_dataset before the issue fixing
![A thaliana Lucia 07 2022 col-0 30m-1h after Infection 22162057 object
id
0](https://github.com/open-mmlab/mmsegmentation/assets/66273343/5431472a-83fd-485f-aeb7-c65f27f1993d)

Example viz_dataset after the issue fixing
![A thaliana Lucia 07 2022 col-0 30m-1h after Infection 22162057 object
id
0](https://github.com/open-mmlab/mmsegmentation/assets/66273343/3d6d4937-41f0-4a18-ae47-35bc43d78843)

---------

Co-authored-by: xiexinch <xiexinch@outlook.com>
---
 mmseg/datasets/transforms/transforms.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/mmseg/datasets/transforms/transforms.py b/mmseg/datasets/transforms/transforms.py
index da3bd5f936..1571c2279c 100644
--- a/mmseg/datasets/transforms/transforms.py
+++ b/mmseg/datasets/transforms/transforms.py
@@ -2206,10 +2206,7 @@ def __init__(self,
         self.aug = Compose([self.albu_builder(t) for t in self.transforms])
 
         if not keymap:
-            self.keymap_to_albu = {
-                'img': 'image',
-                'gt_masks': 'masks',
-            }
+            self.keymap_to_albu = {'img': 'image', 'gt_seg_map': 'mask'}
         else:
             self.keymap_to_albu = copy.deepcopy(keymap)
         self.keymap_back = {v: k for k, v in self.keymap_to_albu.items()}