[Enhancement] Revise file structure and add ut (#1145)

* refactor mmedit/models/utils/

* rename layers to base_archs

* move fid-inception and inception_utils to evaluation/funcitonal

* merge batch process to edit data processor

* move misc to utils

* rename loops.py to gen_loops.py

* move log_processor to runner

* move trans utils to utils

* move some function from metric_utils to mmedit/utils

* add unit test for inception utils.py

* add unit test for average model

* add unit test for base translation model

* add more ut for sagan generator

* add more unit test for biggan generator

* add unit test for io_utils and revise the name of variable

* add more unit test for biggan deep generator

* remove useless ut files and revise some uts

* add unit test for stylegan3

* add more unit test for cyclegan

* add more unit test for stylegan2 module

* skip unit test of CLIP loss on win-cuda env

* omit ops under stylegan2 folder

* fix ut of alpha.py to avoid the randomness

* update base_archs

* update base_archs

* update ut

* fix import error

* add unit test for singan-modules

* add unit test for sn module

* add unit test for stylegan1 and stylegan3

* revise unit test of arcFace

* add unit test for wgan-module

* complete unit test of gen metric

* remove useless unit test files

* remove useless uts and revise ut checking script

* revise omit file list in CI configs

* update dev_scripts

Co-authored-by: zenggyh1900 <[email protected]>

.dev_scripts/README.md

@@ -0,0 +1,11 @@
+## Check UT
+
+Please check your UT by the following scripts:
+
+```python
+cd mmediting/
+python .dev_script/update_ut.py
+```
+
+Then, you will find some redundant UT, missing UT and blank UT.
+Please create UTs according to your package code implementation.

.dev_scripts/update_ut.py

@@ -1,51 +1,67 @@
 import os
 import os.path as osp
 from argparse import ArgumentParser
+from fnmatch import fnmatch
 from glob import glob
 
 from tqdm import tqdm
 
 parser = ArgumentParser()
 parser.add_argument('--src', type=str, default='mmedit')
 parser.add_argument('--dst', type=str, default='tests')
+parser.add_argument(
+    '--exclude',
+    nargs='+',
+    default=[
+        'mmedit/.mim', 'mmedit/registry.py', 'mmedit/version.py',
+        '__pycache__', '__init__', '**/__init__.py', 
+        '**/stylegan3_ops/*', '**/conv2d_gradfix.py', '**/grid_sample_gradfix.py',
+        '**/misc.py', '**/upfirdn2d.py', 
+        '**/all_gather_layer.py', '**/typing.py'
+    ])
 args = parser.parse_args()
 
 
+def check_exclude(fn):
+    for pattern in args.exclude:
+        if fnmatch(fn, pattern):
+            return True
+    return False
+
+
 def update_ut():
 
-    folders = [f for f in os.listdir(args.src) if osp.isdir(f'mmedit/{f}')]
     target_ut = []
     missing_ut = []
     blank_ut = []
 
-    for subf in folders:
-        if subf == '.mim' or subf == '__pycache__':
+    file_list = glob('mmedit/**/*.py', recursive=True)
+
+    for f in tqdm(file_list):
+        if check_exclude(f):
             continue
 
-        file_list = glob(f'mmedit/{subf}/**/*.py', recursive=True)
-
-        for f in tqdm(file_list, desc=f'mmedit/{subf}'):
-            if osp.splitext(osp.basename(f))[0] != '__init__':
-
-                dirname = osp.dirname(f)
-                dirname = dirname.replace('__', '')
-                dirname = dirname.replace('mmedit', 'tests')
-                dirname = dirname.replace('/', '/test_')
-                os.makedirs(dirname, exist_ok=True)
-
-                basename = osp.basename(f)
-                basename = 'test_' + basename
-
-                dst_path = osp.join(dirname, basename)
-                target_ut.append(dst_path)
-                if not osp.exists(dst_path):
-                    missing_ut.append(dst_path)
-                    fp = open(dst_path, 'a')
-                    fp.close()
-                else:
-                    text_lines = open(dst_path, 'r').readlines()
-                    if len(text_lines) <= 3:
-                        blank_ut.append(dst_path)
+        if osp.splitext(osp.basename(f))[0] != '__init__':
+
+            dirname = osp.dirname(f)
+            dirname = dirname.replace('__', '')
+            dirname = dirname.replace('mmedit', 'tests')
+            dirname = dirname.replace('/', '/test_')
+            os.makedirs(dirname, exist_ok=True)
+
+            basename = osp.basename(f)
+            basename = 'test_' + basename
+
+            dst_path = osp.join(dirname, basename)
+            target_ut.append(dst_path)
+            if not osp.exists(dst_path):
+                missing_ut.append(dst_path)
+                fp = open(dst_path, 'a')
+                fp.close()
+            else:
+                text_lines = open(dst_path, 'r').readlines()
+                if len(text_lines) <= 3:
+                    blank_ut.append(dst_path)
 
     existing_ut = glob('tests/test_*/**/*.py', recursive=True)
     additional_ut = list(set(existing_ut) - set(target_ut))

.github/workflows/merge_stage_test.yml

@@ -97,7 +97,7 @@ jobs:
       - name: Run unittests and generate coverage report
         run: |
           coverage run --branch --source mmedit -m pytest tests/
-          coverage xml --omit="**/stylegan3_ops/*,**/conv2d_gradfix.py"
+          coverage xml --omit="**/stylegan3_ops/*,**/conv2d_gradfix.py,**/grid_sample_gradfix.py,**/misc.py,**/upfirdn2d.py,**all_gather_layer.py"
           coverage report -m
       # Only upload coverage report for python3.7 && pytorch1.8.1 cpu
       - name: Upload coverage to Codecov

.github/workflows/pr_stage_test.yml

@@ -45,7 +45,7 @@ jobs:
       - name: Run unittests and generate coverage report
         run: |
           coverage run --branch --source mmedit -m pytest tests/
-          coverage xml --omit="**/stylegan3_ops/*,**/conv2d_gradfix.py"
+          coverage xml --omit="**/stylegan3_ops/*,**/conv2d_gradfix.py,**/grid_sample_gradfix.py,**/misc.py,**/upfirdn2d.py,**all_gather_layer.py"
           coverage report -m
       # Upload coverage report for python3.7 && pytorch1.8.1 cpu
       - name: Upload coverage to Codecov

demo/inpainting_demo.py

@@ -5,7 +5,7 @@
 import torch
 
 from mmedit.apis import init_model, inpainting_inference
-from mmedit.engine import tensor2img
+from mmedit.utils import tensor2img
 
 
 def parse_args():

demo/restoration_demo.py

@@ -6,8 +6,7 @@
 import torch
 
 from mmedit.apis import init_model, restoration_inference
-from mmedit.engine import tensor2img
-from mmedit.utils import modify_args
+from mmedit.utils import modify_args, tensor2img
 
 
 def parse_args():

demo/restoration_video_demo.py

@@ -8,8 +8,7 @@
 import torch
 
 from mmedit.apis import init_model, restoration_video_inference
-from mmedit.engine import tensor2img
-from mmedit.utils import modify_args
+from mmedit.utils import modify_args, tensor2img
 
 VIDEO_EXTENSIONS = ('.mp4', '.mov')
 

mmedit/datasets/transforms/__init__.py

@@ -25,28 +25,24 @@
                                   RandomJPEGCompression, RandomNoise,
                                   RandomResize, RandomVideoCompression)
 from .random_down_sampling import RandomDownSampling
-from .trans_utils import (adjust_gamma, bbox2mask, brush_stroke_mask,
-                          get_irregular_mask, random_bbox)
 from .trimap import (FormatTrimap, GenerateTrimap,
                      GenerateTrimapWithDistTransform, TransformTrimap)
 from .values import CopyValues, SetValues
 
 __all__ = [
-    'random_bbox', 'get_irregular_mask', 'brush_stroke_mask', 'bbox2mask',
-    'adjust_gamma', 'BinarizeImage', 'Clip', 'ColorJitter', 'CopyValues',
-    'Crop', 'CropLike', 'DegradationsWithShuffle', 'LoadImageFromFile',
-    'LoadMask', 'Flip', 'FixedCrop', 'GenerateCoordinateAndCell',
-    'GenerateFacialHeatmap', 'GenerateFrameIndices',
-    'GenerateFrameIndiceswithPadding', 'GenerateSegmentIndices',
-    'GetMaskedImage', 'GetSpatialDiscountMask', 'MATLABLikeResize',
-    'MirrorSequence', 'ModCrop', 'Normalize', 'PackEditInputs',
-    'PairedRandomCrop', 'RandomAffine', 'RandomBlur', 'RandomDownSampling',
-    'RandomJPEGCompression', 'RandomMaskDilation', 'RandomNoise',
-    'RandomResize', 'RandomResizedCrop', 'RandomRotation', 'RandomTransposeHW',
-    'RandomVideoCompression', 'RescaleToZeroOne', 'Resize', 'SetValues',
-    'TemporalReverse', 'ToTensor', 'UnsharpMasking', 'CropAroundCenter',
-    'CropAroundFg', 'GenerateSeg', 'CropAroundUnknown', 'GenerateSoftSeg',
-    'FormatTrimap', 'TransformTrimap', 'GenerateTrimap',
+    'BinarizeImage', 'Clip', 'ColorJitter', 'CopyValues', 'Crop', 'CropLike',
+    'DegradationsWithShuffle', 'LoadImageFromFile', 'LoadMask', 'Flip',
+    'FixedCrop', 'GenerateCoordinateAndCell', 'GenerateFacialHeatmap',
+    'GenerateFrameIndices', 'GenerateFrameIndiceswithPadding',
+    'GenerateSegmentIndices', 'GetMaskedImage', 'GetSpatialDiscountMask',
+    'MATLABLikeResize', 'MirrorSequence', 'ModCrop', 'Normalize',
+    'PackEditInputs', 'PairedRandomCrop', 'RandomAffine', 'RandomBlur',
+    'RandomDownSampling', 'RandomJPEGCompression', 'RandomMaskDilation',
+    'RandomNoise', 'RandomResize', 'RandomResizedCrop', 'RandomRotation',
+    'RandomTransposeHW', 'RandomVideoCompression', 'RescaleToZeroOne',
+    'Resize', 'SetValues', 'TemporalReverse', 'ToTensor', 'UnsharpMasking',
+    'CropAroundCenter', 'CropAroundFg', 'GenerateSeg', 'CropAroundUnknown',
+    'GenerateSoftSeg', 'FormatTrimap', 'TransformTrimap', 'GenerateTrimap',
     'GenerateTrimapWithDistTransform', 'CompositeFg', 'RandomLoadResizeBg',
     'MergeFgAndBg', 'PerturbBg', 'RandomJitter', 'LoadPairedImageFromFile',
     'CenterCropLongEdge', 'RandomCropLongEdge', 'NumpyPad'

mmedit/datasets/transforms/alpha.py

@@ -10,7 +10,7 @@
 from mmengine.utils import is_list_of, is_tuple_of
 
 from mmedit.registry import TRANSFORMS
-from .trans_utils import random_choose_unknown
+from mmedit.utils import random_choose_unknown
 
 
 @TRANSFORMS.register_module()

mmedit/datasets/transforms/crop.py

@@ -9,7 +9,7 @@
 from torch.nn.modules.utils import _pair
 
 from mmedit.registry import TRANSFORMS
-from .trans_utils import random_choose_unknown
+from mmedit.utils import random_choose_unknown
 
 
 @TRANSFORMS.register_module()

mmedit/datasets/transforms/fgbg.py

@@ -10,7 +10,7 @@
 from mmengine.fileio import FileClient
 
 from mmedit.registry import TRANSFORMS
-from .trans_utils import add_gaussian_noise, adjust_gamma
+from mmedit.utils import add_gaussian_noise, adjust_gamma
 
 
 @TRANSFORMS.register_module()

mmedit/datasets/transforms/generate_assistant.py

@@ -4,7 +4,7 @@
 from mmcv.transforms.base import BaseTransform
 
 from mmedit.registry import TRANSFORMS
-from .trans_utils import make_coord
+from mmedit.utils import make_coord
 
 try:
     import face_alignment

mmedit/datasets/transforms/loading.py

@@ -8,7 +8,7 @@
 from mmengine.fileio import FileClient, list_from_file
 
 from mmedit.registry import TRANSFORMS
-from .trans_utils import (bbox2mask, brush_stroke_mask, get_irregular_mask,
+from mmedit.utils import (bbox2mask, brush_stroke_mask, get_irregular_mask,
                           random_bbox)
 
 

mmedit/engine/__init__.py

@@ -1,7 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from .hooks import *  # noqa: F401, F403
-from .logging import *  # noqa: F401, F403
-from .misc import *  # noqa: F401, F403
 from .optimizers import *  # noqa: F401, F403
 from .runner import *  # noqa: F401, F403
 from .schedulers import *  # noqa: F401, F403

mmedit/engine/runner/__init__.py

@@ -1,5 +1,9 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-from .loops import GenTestLoop, GenValLoop
+from .gen_loops import GenTestLoop, GenValLoop
+from .log_processor import GenLogProcessor
 from .multi_loops import MultiTestLoop, MultiValLoop
 
-__all__ = ['MultiValLoop', 'MultiTestLoop', 'GenTestLoop', 'GenValLoop']
+__all__ = [
+    'MultiValLoop', 'MultiTestLoop', 'GenTestLoop', 'GenValLoop',
+    'GenLogProcessor'
+]

mmedit/evaluation/functional/__init__.py

@@ -1,4 +1,10 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+from .fid_inception import InceptionV3
 from .gaussian_funcs import gauss_gradient
+from .inception_utils import (disable_gpu_fuser_on_pt19, load_inception,
+                              prepare_inception_feat, prepare_vgg_feat)
 
-__all__ = ['gauss_gradient']
+__all__ = [
+    'gauss_gradient', 'InceptionV3', 'disable_gpu_fuser_on_pt19',
+    'load_inception', 'prepare_vgg_feat', 'prepare_inception_feat'
+]

mmedit/evaluation/metrics/inception_utils.py → mmedit/evaluation/functional/inception_utils.py

@@ -21,8 +21,8 @@
 from torch.utils.data.dataset import Dataset
 from torchvision.models.inception import inception_v3
 
-from mmedit.models import InceptionV3
-from mmedit.utils import MMGEN_CACHE_DIR, download_from_url
+from mmedit.utils import MMEDIT_CACHE_DIR, download_from_url
+from . import InceptionV3
 
 ALLOWED_INCEPTION = ['StyleGAN', 'PyTorch']
 TERO_INCEPTION_URL = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/inception-2015-12-05.pt'  # noqa
@@ -131,7 +131,7 @@ def _load_inception_from_url(inception_url: str) -> nn.Module:
     print_log(f'Try to download Inception Model from {inception_url}...',
               'current')
     try:
-        path = download_from_url(inception_url, dest_dir=MMGEN_CACHE_DIR)
+        path = download_from_url(inception_url, dest_dir=MMEDIT_CACHE_DIR)
         print_log('Download Finished.', 'current')
         return _load_inception_from_path(path)
     except Exception as e:
@@ -336,7 +336,7 @@ def prepare_inception_feat(dataloader: DataLoader,
     if inception_pkl is None:
         inception_pkl, args = get_inception_feat_cache_name_and_args(
             dataloader, metric, real_nums, capture_mean_cov, capture_all)
-        inception_pkl = osp.join(MMGEN_CACHE_DIR, inception_pkl)
+        inception_pkl = osp.join(MMEDIT_CACHE_DIR, inception_pkl)
     else:
         args = dict()
     if osp.exists(inception_pkl):
@@ -510,7 +510,7 @@ def prepare_vgg_feat(dataloader: DataLoader,
     # cannot load or download from file, extract manually
     if vgg_pkl is None:
         vgg_pkl, args = get_vgg_feat_cache_name_and_args(dataloader, metric)
-        vgg_pkl = osp.join(MMGEN_CACHE_DIR, vgg_pkl)
+        vgg_pkl = osp.join(MMEDIT_CACHE_DIR, vgg_pkl)
     else:
         args = dict()
     if osp.exists(vgg_pkl):

mmedit/evaluation/metrics/fid.py

@@ -10,9 +10,9 @@
 from torch.utils.data.dataloader import DataLoader
 
 from mmedit.registry import METRICS
+from ..functional import (disable_gpu_fuser_on_pt19, load_inception,
+                          prepare_inception_feat)
 from .base_gen_metric import GenerativeMetric
-from .inception_utils import (disable_gpu_fuser_on_pt19, load_inception,
-                              prepare_inception_feat)
 
 
 @METRICS.register_module('FID-Full')

mmedit/evaluation/metrics/inception_score.py

@@ -12,8 +12,9 @@
 from torch.utils.data.dataloader import DataLoader
 
 from mmedit.registry import METRICS
+# from .inception_utils import disable_gpu_fuser_on_pt19, load_inception
+from ..functional import disable_gpu_fuser_on_pt19, load_inception
 from .base_gen_metric import GenerativeMetric
-from .inception_utils import disable_gpu_fuser_on_pt19, load_inception
 
 
 @METRICS.register_module('IS')