[Feature] Support NAFNet model

configs/nafnet/nafnet_c64eb11128mb1db1111_lr1e-3_400k_gopro.py

@@ -0,0 +1,133 @@
+_base_ = '../_base_/default_runtime.py'
+
+experiment_name = 'nafnet_c64eb11128mb1db1111_lr1e-3_400k_gopro'
+work_dir = f'./work_dirs/{experiment_name}'
+save_dir = './work_dirs/'
+
+# DistributedDataParallel
+model_wrapper_cfg = dict(type='MMSeparateDistributedDataParallel')
+
+# model settings
+model = dict(
+    type='BaseEditModel',
+    generator=dict(
+        type='NAFNetLocal',
+        img_channel=3,
+        mid_channels=64,
+        enc_blk_nums=[1, 1, 1, 28],
+        middle_blk_num=1,
+        dec_blk_nums=[1, 1, 1, 1],
+    ),
+    pixel_loss=dict(type='PSNRLoss'),
+    train_cfg=dict(),
+    test_cfg=dict(),
+    data_preprocessor=dict(
+        type='EditDataPreprocessor',
+        mean=[0., 0., 0.],
+        std=[255., 255., 255.],
+    )
+)
+
+train_pipeline = [
+    dict(type='LoadImageFromFile', key='img'),
+    dict(type='LoadImageFromFile', key='gt'),
+    dict(type='SetValues', dictionary=dict(scale=1)),
+    dict(
+        type='Flip',
+        keys=['img', 'gt'],
+        flip_ratio=0.5,
+        direction='horizontal'),
+    dict(
+        type='Flip', 
+        keys=['img', 'gt'], 
+        flip_ratio=0.5, 
+        direction='vertical'),
+    dict(type='RandomTransposeHW', keys=['img', 'gt'], transpose_ratio=0.5),
+    dict(type='PairedRandomCrop', gt_patch_size=256),
+    dict(type='PackEditInputs')
+]
+
+val_pipeline = [
+    dict(type='LoadImageFromFile', key='img'),
+    dict(type='LoadImageFromFile', key='gt'),
+    dict(type='PackEditInputs')
+]
+
+# dataset settings
+dataset_type = 'BasicImageDataset'
+
+train_dataloader = dict(
+    num_workers=8,
+    batch_size=8,  # gpus 4
+    persistent_workers=False,
+    sampler=dict(type='InfiniteSampler', shuffle=True),
+    dataset=dict(
+        type=dataset_type,
+        metainfo=dict(dataset_type='gopro', task_name='deblur'),
+        data_root='../datasets/gopro/train',
+        data_prefix=dict(gt='sharp', img='blur'),
+        ann_file='meta_info_gopro_train.txt',
+        pipeline=train_pipeline))
+
+val_dataloader = dict(
+    num_workers=4,
+    persistent_workers=False,
+    drop_last=False,
+    sampler=dict(type='DefaultSampler', shuffle=False),
+    dataset=dict(
+        type=dataset_type,
+        metainfo=dict(dataset_type='gopro', task_name='deblur'),
+        data_root='../datasets/gopro/test',
+        ann_file='meta_info_gopro_test.txt',
+        data_prefix=dict(gt='sharp', img='blur'),
+        pipeline=val_pipeline))
+
+test_dataloader = val_dataloader
+
+val_evaluator = [
+    dict(type='MAE'),
+    dict(type='PSNR'),
+    dict(type='SSIM'),
+]
+test_evaluator = val_evaluator
+
+train_cfg = dict(
+    type='IterBasedTrainLoop', max_iters=400_000, val_interval=20000)
+val_cfg = dict(type='ValLoop')
+test_cfg = dict(type='TestLoop')
+
+# optimizer
+# optim_wrapper = dict(
+#     constructor='MultiOptimWrapperConstructor',
+#     generator=dict(
+#         type='OptimWrapper',
+#         optimizer=dict(type='Adam', lr=1e-3, betas=(0.9, 0.9))))
+optim_wrapper = dict(
+    constructor='DefaultOptimWrapperConstructor',
+    type='OptimWrapper',
+    optimizer=dict(type='AdamW', lr=1e-3, weight_decay=1e-3, betas=(0.9, 0.9)))
+
+# learning policy
+param_scheduler = dict(
+    type='CosineAnnealingLR',
+    by_epoch=False,
+    T_max=400_000, 
+    eta_min=1e-7)
+
+default_hooks = dict(
+    checkpoint=dict(
+        type='CheckpointHook',
+        interval=5000,
+        save_optimizer=True,
+        by_epoch=False,
+        out_dir=save_dir,
+    ),
+    timer=dict(type='IterTimerHook'),
+    logger=dict(type='LoggerHook', interval=100),
+    param_scheduler=dict(type='ParamSchedulerHook'),
+    sampler_seed=dict(type='DistSamplerSeedHook'),
+)
+
+visualizer = dict(bgr2rgb=False)
+
+randomness = dict(seed=10, diff_rank_seed=True)

configs/nafnet/nafnet_c64eb2248mb12db2222_lr1e-3_400k_sidd.py

@@ -0,0 +1,127 @@
+_base_ = '../_base_/default_runtime.py'
+
+experiment_name = 'nafnet_c64eb2248mb12db2222_lr1e-3_400k_sidd'
+work_dir = f'./work_dirs/{experiment_name}'
+save_dir = './work_dirs/'
+
+# DistributedDataParallel
+model_wrapper_cfg = dict(type='MMSeparateDistributedDataParallel')
+
+# model settings
+model = dict(
+    type='BaseEditModel',
+    generator=dict(
+        type='NAFNet',
+        img_channel=3,
+        mid_channels=64,
+        enc_blk_nums=[2, 2, 4, 8],
+        middle_blk_num=12,
+        dec_blk_nums=[2, 2, 2, 2],
+    ),
+    pixel_loss=dict(type='PSNRLoss'),
+    train_cfg=dict(),
+    test_cfg=dict(),
+    data_preprocessor=dict(
+        type='EditDataPreprocessor',
+        mean=[0.0, 0.0, 0.0],
+        std=[255.0, 255.0, 255.0],
+    )
+)
+
+train_pipeline = [
+    dict(type='LoadImageFromFile', key='img'),
+    dict(type='LoadImageFromFile', key='gt'),
+    dict(type='SetValues', dictionary=dict(scale=1)),
+    dict(
+        type='Flip',
+        keys=['img', 'gt'],
+        flip_ratio=0.5,
+        direction='horizontal'),
+    dict(
+        type='Flip', 
+        keys=['img', 'gt'], 
+        flip_ratio=0.5, 
+        direction='vertical'),
+    dict(type='RandomTransposeHW', keys=['img', 'gt'], transpose_ratio=0.5),
+    dict(type='PairedRandomCrop', gt_patch_size=256),
+    dict(type='PackEditInputs')
+]
+
+val_pipeline = [
+    dict(type='LoadImageFromFile', key='img'),
+    dict(type='LoadImageFromFile', key='gt'),
+    dict(type='PackEditInputs')
+]
+
+# dataset settings
+dataset_type = 'BasicImageDataset'
+
+train_dataloader = dict(
+    num_workers=8,
+    batch_size=8,  # gpus 4
+    persistent_workers=False,
+    sampler=dict(type='InfiniteSampler', shuffle=True),
+    dataset=dict(
+        type=dataset_type,
+        metainfo=dict(dataset_type='sidd', task_name='denoising'),
+        data_root='../datasets/SIDD/train',
+        data_prefix=dict(gt='gt', img='noisy'),
+        filename_tmpl=dict(img='{}_NOISY', gt='{}_GT'), 
+        pipeline=train_pipeline))
+
+val_dataloader = dict(
+    num_workers=4,
+    persistent_workers=False,
+    drop_last=False,
+    sampler=dict(type='DefaultSampler', shuffle=False),
+    dataset=dict(
+        type=dataset_type,
+        metainfo=dict(dataset_type='sidd', task_name='denoising'),
+        data_root='../datasets/SIDD/val/',
+        data_prefix=dict(gt='gt', img='noisy'),
+        filename_tmpl=dict(gt='{}_GT', img='{}_NOISY'), 
+        pipeline=val_pipeline))
+
+test_dataloader = val_dataloader
+
+val_evaluator = [
+    dict(type='MAE'),
+    dict(type='PSNR'),
+    dict(type='SSIM'),
+]
+test_evaluator = val_evaluator
+
+train_cfg = dict(
+    type='IterBasedTrainLoop', max_iters=400_000, val_interval=20000)
+val_cfg = dict(type='ValLoop')
+test_cfg = dict(type='TestLoop')
+
+# optimizer
+optim_wrapper = dict(
+    constructor='MultiOptimWrapperConstructor',
+    generator=dict(
+        type='OptimWrapper',
+        optimizer=dict(type='AdamW', lr=1e-3, betas=(0.9, 0.9))))
+
+# learning policy
+param_scheduler = dict(
+    type='CosineAnnealingLR',
+    by_epoch=False,
+    T_max=400_000, 
+    eta_min=1e-7)
+
+default_hooks = dict(
+    checkpoint=dict(
+        type='CheckpointHook',
+        interval=5000,
+        save_optimizer=True,
+        by_epoch=False,
+        out_dir=save_dir,
+    ),
+    timer=dict(type='IterTimerHook'),
+    logger=dict(type='LoggerHook', interval=100),
+    param_scheduler=dict(type='ParamSchedulerHook'),
+    sampler_seed=dict(type='DistSamplerSeedHook'),
+)
+
+visualizer = dict(bgr2rgb=True)

configs/real_esrgan/realesrnet_c64b23g32_4xb12-lr2e-4-1000k_df2k-ost.py

@@ -261,4 +261,4 @@
         interval=1,
         interp_cfg=dict(momentum=0.999),
     )
-]
+]

mmedit/evaluation/metrics/ssim.py

@@ -180,3 +180,60 @@ def ssim(img1,
         ssims.append(_ssim(img1[..., i], img2[..., i]))
 
     return np.array(ssims).mean()
+
+# Components for SSIM with 3D kernels
+# Not used for now
+import torch
+
+def _3d_gaussian_calculator(img, conv3d):
+    out = conv3d(img.unsqueeze(0).unsqueeze(0)).squeeze(0).squeeze(0)
+    return out
+
+def _generate_3d_gaussian_kernel():
+    kernel = cv2.getGaussianKernel(11, 1.5)
+    window = np.outer(kernel, kernel.transpose())
+    kernel_3 = cv2.getGaussianKernel(11, 1.5)
+    kernel = torch.tensor(np.stack([window * k for k in kernel_3], axis=0))
+    conv3d = torch.nn.Conv3d(1, 1, (11, 11, 11), stride=1, padding=(5, 5, 5), bias=False, padding_mode='replicate')
+    conv3d.weight.requires_grad = False
+    conv3d.weight[0, 0, :, :, :] = kernel
+    return conv3d
+
+def _ssim_3d(img1, img2, max_value):
+    assert len(img1.shape) == 3 and len(img2.shape) == 3
+    """Calculate SSIM (structural similarity) for one channel images.
+
+    It is called by func:`calculate_ssim`.
+
+    Args:
+        img1 (ndarray): Images with range [0, 255]/[0, 1] with order 'HWC'.
+        img2 (ndarray): Images with range [0, 255]/[0, 1] with order 'HWC'.
+
+    Returns:
+        float: ssim result.
+    """
+    C1 = (0.01 * max_value) ** 2
+    C2 = (0.03 * max_value) ** 2
+    img1 = img1.astype(np.float64)
+    img2 = img2.astype(np.float64)
+
+    kernel = _generate_3d_gaussian_kernel().cuda()
+
+    img1 = torch.tensor(img1).float().cuda()
+    img2 = torch.tensor(img2).float().cuda()
+
+
+    mu1 = _3d_gaussian_calculator(img1, kernel)
+    mu2 = _3d_gaussian_calculator(img2, kernel)
+
+    mu1_sq = mu1 ** 2
+    mu2_sq = mu2 ** 2
+    mu1_mu2 = mu1 * mu2
+    sigma1_sq = _3d_gaussian_calculator(img1 ** 2, kernel) - mu1_sq
+    sigma2_sq = _3d_gaussian_calculator(img2 ** 2, kernel) - mu2_sq
+    sigma12 = _3d_gaussian_calculator(img1*img2, kernel) - mu1_mu2
+
+    ssim_map = ((2 * mu1_mu2 + C1) *
+                (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
+                                       (sigma1_sq + sigma2_sq + C2))
+    return float(ssim_map.mean())

mmedit/models/editors/__init__.py

@@ -50,6 +50,7 @@
 from .tof import TOFlowVFINet, TOFlowVSRNet, ToFResBlock
 from .ttsr import LTE, TTSR, SearchTransformer, TTSRDiscriminator, TTSRNet
 from .wgan_gp import WGANGP
+from .nafnet import NAFNet, NAFNetLocal, Baseline, BaselineLocal
 
 __all__ = [
     'AOTEncoderDecoder', 'AOTBlockNeck', 'AOTInpaintor',

mmedit/models/editors/nafnet/__init__.py

@@ -0,0 +1,9 @@
+from .nafnet_net import NAFNet, NAFNetLocal
+from .baseline_net import Baseline, BaselineLocal
+
+__all__ = [
+    "NAFNet",
+    "NAFNetLocal",
+    "Baseline",
+    "BaselineLocal",
+]