From 740aa78c30d99660b1a129f6e7b8cd3ca6c38b3b Mon Sep 17 00:00:00 2001 From: liyinshuo Date: Mon, 28 Jun 2021 15:13:58 +0800 Subject: [PATCH 1/2] Update --- tests/test_models/test_restorers/test_ttsr.py | 21 ++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/tests/test_models/test_restorers/test_ttsr.py b/tests/test_models/test_restorers/test_ttsr.py index f0758b2df9..637e3188bd 100644 --- a/tests/test_models/test_restorers/test_ttsr.py +++ b/tests/test_models/test_restorers/test_ttsr.py @@ -72,7 +72,26 @@ def test_ttsr(): num_blocks=(16, 16, 8, 4)), extractor=dict(type='LTE'), transformer=dict(type='SearchTransformer'), - pixel_loss=dict(type='L1Loss', loss_weight=1.0, reduction='mean')) + discriminator=dict(type='TTSRDiscriminator', in_size=64), + pixel_loss=dict(type='L1Loss', loss_weight=1.0, reduction='mean'), + perceptual_loss=dict( + type='PerceptualLoss', + layer_weights={'29': 1.0}, + vgg_type='vgg19', + perceptual_weight=1e-2, + style_weight=0.001, + criterion='mse'), + transferal_perceptual_loss=dict( + type='TransferalPerceptualLoss', + loss_weight=1e-2, + use_attention=False, + criterion='mse'), + gan_loss=dict( + type='GANLoss', + gan_type='vanilla', + loss_weight=1e-3, + real_label_val=1.0, + fake_label_val=0)) scale = 4 train_cfg = None From 885071ad5c6768b0f0add03bc5bdbf2b597d1625 Mon Sep 17 00:00:00 2001 From: liyinshuo Date: Mon, 28 Jun 2021 19:52:49 +0800 Subject: [PATCH 2/2] [Feature] Add config of ttsr-gan --- configs/restorers/ttsr/README.md | 7 +- .../ttsr/ttsr-gan_x4_c64b16_g1_500k_CUFED.py | 274 ++++++++++++++++++ 2 files changed, 278 insertions(+), 3 deletions(-) create mode 100644 configs/restorers/ttsr/ttsr-gan_x4_c64b16_g1_500k_CUFED.py diff --git a/configs/restorers/ttsr/README.md b/configs/restorers/ttsr/README.md index ffc3b9ac60..0bcb9b51bc 100644 --- a/configs/restorers/ttsr/README.md +++ b/configs/restorers/ttsr/README.md @@ -20,6 +20,7 @@ Evaluated on RGB channels, `scale` pixels in each border are cropped before eval The metrics are `PSNR / SSIM`. -| Method | scale | CUFED | Download | -| :---------------------------------------------------------------------------------------------: | :---: | :--------------: | :-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------: | -| [ttsr-rec_x4_c64b16_g1_200k_CUFED](/configs/restorers/ttsr/ttsr-rec_x4_c64b16_g1_200k_CUFED.py) | x4 | 25.2433 / 0.7491 | [model](https://download.openmmlab.com/mmediting/restorers/ttsr/ttsr-rec_x4_c64b16_g1_200k_CUFED_20210525-b0dba584.pth?versionId=CAEQKxiBgIDht5ONzRciIDdjZTQ1NmFmYzhjNjQ5NGFhNjkyNzU1N2UxMjMyZWE4) \| [log](https://download.openmmlab.com/mmediting/restorers/ttsr/ttsr-rec_x4_c64b16_g1_200k_CUFED_20210525-b0dba584.log.json?versionId=CAEQKxiCgMCnuJONzRciIDUzNmVkNGNmNTlkMDQzMmFhZDAzYzQ5NmUzNTI5YmYz) | +| Method | scale | CUFED | Download | +| :---------------------------------------------------------------------------------------------: | :---: | :--------------: | :---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------: | +| [ttsr-rec_x4_c64b16_g1_200k_CUFED](/configs/restorers/ttsr/ttsr-rec_x4_c64b16_g1_200k_CUFED.py) | x4 | 25.2433 / 0.7491 | [model](https://download.openmmlab.com/mmediting/restorers/ttsr/ttsr-rec_x4_c64b16_g1_200k_CUFED_20210525-b0dba584.pth) \| [log](https://download.openmmlab.com/mmediting/restorers/ttsr/ttsr-rec_x4_c64b16_g1_200k_CUFED_20210525-b0dba584.log.json) | +| [ttsr-gan_x4_c64b16_g1_500k_CUFED](/configs/restorers/ttsr/ttsr-gan_x4_c64b16_g1_500k_CUFED.py) | x4 | 24.6075 / 0.7234 | [model](https://download.openmmlab.com/mmediting/restorers/ttsr/ttsr-gan_x4_c64b16_g1_500k_CUFED_20210626-2ab28ca0.pth) \| [log](https://download.openmmlab.com/mmediting/restorers/ttsr/ttsr-gan_x4_c64b16_g1_500k_CUFED_20210626-2ab28ca0.log.json) | diff --git a/configs/restorers/ttsr/ttsr-gan_x4_c64b16_g1_500k_CUFED.py b/configs/restorers/ttsr/ttsr-gan_x4_c64b16_g1_500k_CUFED.py new file mode 100644 index 0000000000..03fb8d7223 --- /dev/null +++ b/configs/restorers/ttsr/ttsr-gan_x4_c64b16_g1_500k_CUFED.py @@ -0,0 +1,274 @@ +exp_name = 'ttsr-gan_x4_c64b16_g1_500k_CUFED' +scale = 4 + +# model settings +model = dict( + type='TTSR', + generator=dict( + type='TTSRNet', + in_channels=3, + out_channels=3, + mid_channels=64, + num_blocks=(16, 16, 8, 4)), + extractor=dict(type='LTE'), + transformer=dict(type='SearchTransformer'), + discriminator=dict(type='TTSRDiscriminator', in_size=160), + pixel_loss=dict(type='L1Loss', loss_weight=1.0, reduction='mean'), + perceptual_loss=dict( + type='PerceptualLoss', + layer_weights={'29': 1.0}, + vgg_type='vgg19', + perceptual_weight=1e-2, + style_weight=0, + criterion='mse'), + transferal_perceptual_loss=dict( + type='TransferalPerceptualLoss', + loss_weight=1e-2, + use_attention=False, + criterion='mse'), + gan_loss=dict( + type='GANLoss', + gan_type='vanilla', + loss_weight=1e-3, + real_label_val=1.0, + fake_label_val=0)) +# model training and testing settings +train_cfg = dict(fix_iter=25000, disc_steps=2) +test_cfg = dict(metrics=['PSNR', 'SSIM'], crop_border=scale) + +# dataset settings +train_dataset_type = 'SRFolderRefDataset' +val_dataset_type = 'SRFolderRefDataset' +test_dataset_type = 'SRFolderRefDataset' +train_pipeline = [ + dict( + type='LoadImageFromFile', + io_backend='disk', + key='gt', + flag='color', + channel_order='rgb', + backend='pillow'), + dict( + type='LoadImageFromFile', + io_backend='disk', + key='ref', + flag='color', + channel_order='rgb', + backend='pillow'), + dict(type='CropLike', target_key='ref', reference_key='gt'), + dict( + type='Resize', + scale=1 / scale, + keep_ratio=True, + keys=['gt', 'ref'], + output_keys=['lq', 'ref_down'], + interpolation='bicubic', + backend='pillow'), + dict( + type='Resize', + scale=float(scale), + keep_ratio=True, + keys=['lq', 'ref_down'], + output_keys=['lq_up', 'ref_downup'], + interpolation='bicubic', + backend='pillow'), + dict( + type='Normalize', + keys=['lq', 'gt'], + mean=[127.5, 127.5, 127.5], + std=[127.5, 127.5, 127.5]), + dict( + type='Normalize', + keys=['lq_up', 'ref', 'ref_downup'], + mean=[0., 0., 0.], + std=[255., 255., 255.]), + dict( + type='Flip', + keys=['lq', 'gt', 'lq_up'], + flip_ratio=0.5, + direction='horizontal'), + dict( + type='Flip', + keys=['lq', 'gt', 'lq_up'], + flip_ratio=0.5, + direction='vertical'), + dict( + type='RandomTransposeHW', + keys=['lq', 'gt', 'lq_up'], + transpose_ratio=0.5), + dict( + type='Flip', + keys=['ref', 'ref_downup'], + flip_ratio=0.5, + direction='horizontal'), + dict( + type='Flip', + keys=['ref', 'ref_downup'], + flip_ratio=0.5, + direction='vertical'), + dict( + type='RandomTransposeHW', + keys=['ref', 'ref_downup'], + transpose_ratio=0.5), + dict( + type='ImageToTensor', keys=['lq', 'gt', 'lq_up', 'ref', 'ref_downup']), + dict( + type='Collect', + keys=['lq', 'gt', 'lq_up', 'ref', 'ref_downup'], + meta_keys=['gt_path', 'ref_path']) +] +valid_pipeline = [ + dict( + type='LoadImageFromFile', + io_backend='disk', + key='gt', + flag='color', + channel_order='rgb', + backend='pillow'), + dict( + type='LoadImageFromFile', + io_backend='disk', + key='ref', + flag='color', + channel_order='rgb', + backend='pillow'), + dict(type='CropLike', target_key='ref', reference_key='gt'), + dict( + type='Resize', + scale=1 / scale, + keep_ratio=True, + keys=['gt', 'ref'], + output_keys=['lq', 'ref_down'], + interpolation='bicubic', + backend='pillow'), + dict( + type='Resize', + scale=float(scale), + keep_ratio=True, + keys=['lq', 'ref_down'], + output_keys=['lq_up', 'ref_downup'], + interpolation='bicubic', + backend='pillow'), + dict( + type='Normalize', + keys=['lq', 'gt'], + mean=[127.5, 127.5, 127.5], + std=[127.5, 127.5, 127.5]), + dict( + type='Normalize', + keys=['lq_up', 'ref', 'ref_downup'], + mean=[0., 0., 0.], + std=[255., 255., 255.]), + dict( + type='ImageToTensor', keys=['lq', 'gt', 'lq_up', 'ref', 'ref_downup']), + dict( + type='Collect', + keys=['lq', 'gt', 'lq_up', 'ref', 'ref_downup'], + meta_keys=['gt_path', 'ref_path']) +] +test_pipeline = [ + dict( + type='LoadImageFromFile', + io_backend='disk', + key='lq', + flag='color', + channel_order='rgb', + backend='pillow'), + dict( + type='LoadImageFromFile', + io_backend='disk', + key='ref', + flag='color', + channel_order='rgb', + backend='pillow'), + dict( + type='Resize', + scale=1 / scale, + keep_ratio=True, + keys=['ref'], + output_keys=['ref_down'], + interpolation='bicubic', + backend='pillow'), + dict( + type='Resize', + scale=float(scale), + keep_ratio=True, + keys=['lq', 'ref_down'], + output_keys=['lq_up', 'ref_downup'], + interpolation='bicubic', + backend='pillow'), + dict( + type='Normalize', + keys=['lq'], + mean=[127.5, 127.5, 127.5], + std=[127.5, 127.5, 127.5]), + dict( + type='Normalize', + keys=['lq_up', 'ref', 'ref_downup'], + mean=[0., 0., 0.], + std=[255., 255., 255.]), + dict(type='ImageToTensor', keys=['lq', 'lq_up', 'ref', 'ref_downup']), + dict( + type='Collect', + keys=['lq', 'lq_up', 'ref', 'ref_downup'], + meta_keys=['lq_path', 'ref_path']) +] + +data = dict( + workers_per_gpu=9, + train_dataloader=dict(samples_per_gpu=9, drop_last=True), + val_dataloader=dict(samples_per_gpu=1), + test_dataloader=dict(samples_per_gpu=1), + train=dict( + type='RepeatDataset', + times=52, + dataset=dict( + type=train_dataset_type, + gt_folder='data/CUFED/train/input/', + ref_folder='data/CUFED/train/ref/', + pipeline=train_pipeline, + scale=scale)), + val=dict( + type=val_dataset_type, + gt_folder='data/CUFED/valid/input_format/', + ref_folder='data/CUFED/valid/ref1_format/', + pipeline=valid_pipeline, + scale=scale), + test=dict( + type=test_dataset_type, + gt_folder='data/CUFED/valid/input_format/', + ref_folder='data/CUFED/valid/ref1_format/', + pipeline=valid_pipeline, + scale=scale)) + +# optimizer +optimizers = dict( + generator=dict(type='Adam', lr=1e-4, betas=(0.9, 0.999)), + discriminator=dict(type='Adam', lr=1e-4, betas=(0.9, 0.999))) + +# learning policy +total_iters = 500000 +lr_config = dict( + policy='Step', + by_epoch=False, + step=[100000, 200000, 300000, 400000], + gamma=0.5) + +checkpoint_config = dict(interval=100, save_optimizer=True, by_epoch=False) +evaluation = dict(interval=5000, save_image=True, gpu_collect=True) +log_config = dict( + interval=100, + hooks=[ + dict(type='TextLoggerHook', by_epoch=False), + # dict(type='TensorboardLoggerHook') + ]) +visual_config = None + +# runtime settings +dist_params = dict(backend='nccl') +log_level = 'INFO' +work_dir = f'./work_dirs/{exp_name}' +load_from = None +resume_from = None +workflow = [('train', 1)] +find_unused_parameters = True