[Refact] support all models in translation inferencer (#1650)

* fix translation load bug. * fix lint. * support cyclegan inference. * fix lint.
open-mmlab · Feb 21, 2023 · a17e934 · a17e934
1 parent 46ae779
commit a17e934
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/mmedit/apis/inferencers/mmedit_inferencer.py b/mmedit/apis/inferencers/mmedit_inferencer.py
@@ -52,7 +52,7 @@ def __init__(self,
         elif self.task in ['inpainting', 'Inpainting']:
             self.inferencer = InpaintingInferencer(
                 config, ckpt, device, extra_parameters, seed=seed)
-        elif self.task in ['translation', 'Image2Image Translation']:
+        elif self.task in ['translation', 'Image2Image']:
             self.inferencer = TranslationInferencer(
                 config, ckpt, device, extra_parameters, seed=seed)
         elif self.task in ['restoration', 'Image Super-Resolution']:

diff --git a/mmedit/apis/inferencers/translation_inferencer.py b/mmedit/apis/inferencers/translation_inferencer.py
@@ -42,16 +42,15 @@ def preprocess(self, img: InputsType) -> Dict:
         test_pipeline = Compose(cfg.test_pipeline)
 
         # prepare data
-        data = dict()
         # dirty code to deal with test data pipeline
+        data = dict()
         data['pair_path'] = img
-        data[f'img_{source_domain}_path'] = img
-        data[f'img_{self.target_domain}_path'] = img
-
+        data['img_A_path'] = img
+        data['img_B_path'] = img
         data = collate([test_pipeline(data)])
         data = self.model.data_preprocessor(data, False)
-        inputs_dict = data['inputs']
 
+        inputs_dict = data['inputs']
         results = inputs_dict[f'img_{source_domain}']
         return results
 

diff --git a/mmedit/edit.py b/mmedit/edit.py
@@ -57,6 +57,7 @@ class MMEdit:
 
         # translation models
         'pix2pix',
+        'cyclegan',
 
         # restoration models
         'esrgan',