Skip to content

Commit

Permalink
[Fix] Avoid mapping unnecessary key to albu format (#1773)
Browse files Browse the repository at this point in the history
  • Loading branch information
liqikai9 authored Oct 27, 2022
1 parent d292bd6 commit 33ecde2
Showing 1 changed file with 11 additions and 22 deletions.
33 changes: 11 additions & 22 deletions mmpose/datasets/transforms/common_transforms.py
Original file line number Diff line number Diff line change
Expand Up @@ -618,7 +618,6 @@ def __init__(self,
}
else:
self.keymap_to_albu = keymap
self.keymap_back = {v: k for k, v in self.keymap_to_albu.items()}

def albu_builder(self, cfg: dict) -> albumentations:
"""Import a module from albumentations.
Expand Down Expand Up @@ -654,23 +653,6 @@ def albu_builder(self, cfg: dict) -> albumentations:

return obj_cls(**args)

@staticmethod
def mapper(d: dict, keymap: dict) -> dict:
"""Dictionary mapper.
Renames keys according to keymap provided.
Args:
d (dict): old dict
keymap (dict): key mapping like {'old_key': 'new_key'}.
Returns:
dict: new dict.
"""

updated_dict = {keymap.get(k, k): v for k, v in d.items()}
return updated_dict

def transform(self, results: dict) -> dict:
"""The transform function of :class:`Albumentation` to apply
albumentations transforms.
Expand All @@ -684,11 +666,18 @@ def transform(self, results: dict) -> dict:
dict: updated result dict.
"""
# map result dict to albumentations format
results = self.mapper(results, self.keymap_to_albu)
results_albu = {}
for k, v in self.keymap_to_albu.items():
assert k in results, \
f'The `{k}` is required to perform albumentations transforms'
results_albu[v] = results[k]

# Apply albumentations transforms
results = self.aug(**results)
# map result dict back to the original format
results = self.mapper(results, self.keymap_back)
results_albu = self.aug(**results_albu)

# map the albu results back to the original format
for k, v in self.keymap_to_albu.items():
results[k] = results_albu[v]

return results

Expand Down

0 comments on commit 33ecde2

Please # to comment.