【图像去噪】第六期论文复现赛——MIRNet

!cd data && tar -xf data140841/SIDD_patches.tar.gz

!unzip data/data150163/pretrained_models.zip -d work/pretrained_models

MIRNet_Paddle    |-- configs                       # 单机单卡/四卡训练配置文件
    |-- dataloaders                # 数据集相关文件
    |-- SIDD_patches
         |-- train                      # SIDD-Medium 训练数据
         |-- val                         # SIDD 测试数据
    |-- networks
         |-- MIRNet_model.py          # MIRNetV1模型代码
         |-- MIRNet_V2_model.py    # MIRNetV2模型代码
    |-- pretrained_models              # 预训练模型
    |-- utils                                     # 一些工具代码
    |-- config.py                             # 配置文件
    |-- losses.py                             # 损失函数
    |-- test_denoising_sidd.py       # 测试SIDD数据上的指标
    |-- train_denoising.py              # 训练代码

!pip install scikit-image natsort yacs

class MSRB(nn.Layer):
    def __init__(self, n_feat, height, width, stride, bias):
        super(MSRB, self).__init__()

        self.n_feat, self.height, self.width = n_feat, height, width
        self.blocks = nn.LayerList([nn.LayerList([DAU(int(n_feat*stride**i))]*width) for i in range(height)])

        INDEX = np.arange(0,width, 2)
        FEATS = [int((stride**i)*n_feat) for i in range(height)]
        SCALE = [2**i for i in range(1,height)]

        self.last_up   = nn.LayerDict()        for i in range(1,height):
            self.last_up.update({f'{i}': UpSample(int(n_feat*stride**i),2**i,stride)})

        self.down = nn.LayerDict()
        self.up   = nn.LayerDict()

        i=0
        SCALE.reverse()        for feat in FEATS:            for scale in SCALE[i:]:
                self.down.update({f'{feat}_{scale}': DownSample(feat,scale,stride)})
            i+=1

        i=0
        FEATS.reverse()        for feat in FEATS:            for scale in SCALE[i:]:                
                self.up.update({f'{feat}_{scale}': UpSample(feat,scale,stride)})
            i+=1

        self.conv_out = nn.Conv2D(n_feat, n_feat, kernel_size=3, padding=1, bias_attr=bias)

        self.selective_kernel = nn.LayerList([SKFF(n_feat*stride**i, height) for i in range(height)])    def forward(self, x):
        inp = x.clone()        #col 1 only
        blocks_out = []        for j in range(self.height):            if j==0:
                inp = self.blocks[j][0](inp)            else:
                inp = self.blocks[j][0](self.down[f'{inp.shape[1]}_{2}'](inp))
            blocks_out.append(inp)        #rest of grid
        for i in range(1,self.width):            #Mesh
            # Replace condition(i%2!=0) with True(Mesh) or False(Plain)
            # if i%2!=0:
            tmp=[]            for j in range(self.height):
                TENSOR = []
                nfeats = (2**j)*self.n_feat                for k in range(self.height):
                    TENSOR.append(self.select_up_down(blocks_out[k], j, k))

                selective_kernel_fusion = self.selective_kernel[j](TENSOR)
                tmp.append(selective_kernel_fusion)            #Forward through either mesh or plain
            for j in range(self.height):
                blocks_out[j] = self.blocks[j][i](tmp[j])        #Sum after grid
        out=[]        for k in range(self.height):
            out.append(self.select_last_up(blocks_out[k], k))  

        out = self.selective_kernel[0](out)

        out = self.conv_out(out)
        out = out + x        return out    def select_up_down(self, tensor, j, k):
        if j==k:            return tensor        else:
            diff = 2 ** np.abs(j-k)            if j<k:                return self.up[f'{tensor.shape[1]}_{diff}'](tensor)            else:                return self.down[f'{tensor.shape[1]}_{diff}'](tensor)    def select_last_up(self, tensor, k):
        if k==0:            return tensor        else:            return self.last_up[f'{k}'](tensor)

class CharbonnierLoss(nn.Layer):
    """Charbonnier Loss (L1)"""

    def __init__(self, eps=1e-3):
        super(CharbonnierLoss, self).__init__()
        self.eps = eps    def forward(self, x, y):
        diff = x - y        # loss = paddle.sum(paddle.sqrt(diff * diff + self.eps))
        loss = paddle.mean(paddle.sqrt((diff * diff) + (self.eps*self.eps)))        return loss

OPTIM:
  BATCH_SIZE: [4]  NUM_EPOCHS: [60]  LR_INITIAL: 2e-4

OPTIM:
  BATCH_SIZE: [8,6,4,2]  NUM_EPOCHS: [30,15,10,5]  LR_INITIAL: 2e-4
  TRAINING:
  PATCH_SIZE: [128,160,192,256]  NUM_WORKERS: [4,4,4,4]

!cd work && python train_denoising.py --model MIRNet --gpus 1

## MIRnet V1# 单机四卡!cd work && python -m paddle.distributed.launch train_denoising.py --model MIRNet --gpus 4## MIRNet V2# 单机单卡!cd work && python train_denoising.py --model MIRNetV2 --gpus 1# 单机四卡!cd work && python -m paddle.distributed.launch train_denoising.py --model MIRNetV2 --gpus 4

# MIRNet V1!cd work && python test_denoising_sidd.py --input_dir ../data/SIDD_patches/val --weights ./pretrained_models/MIRNetV1_paddle.pdparams --model MIRNet

# MIRNet V1PSNR: 39.6872 
SSIM: 0.9586# MIRNet V2PSNR: 39.5286
SSIM: 0.9578

# MIRNet V1!cd work && python predict.py --data_path ./SIDD_patches/val_mini/ --save_path results/ --save_images --model_ckpt ./pretrained_models/MIRNetV1_paddle.pdparams --model MIRNet

# 先上传一张图片，import os.path as ospfrom IPython.display import displayfrom PIL import Image
img_path = 'bird.png' # 改成自己上传的图片名称full_img_path = osp.join(osp.abspath('work/test_images/'), img_path)
img = Image.open(full_img_path).convert('RGB')print('以下为上传的图片:')
display(img)

# MIRNet V1 仅给定干净图片，噪声水平为15!cd work && python predict_single.py --clean_img $full_img_path --save_images --model_ckpt ./pretrained_models/MIRNetV1_paddle.pdparams --model MIRNet

# 去噪效果查看import globfrom IPython.display import displayfrom PIL import Image

imgs = glob.glob('work/test_images/*')for path in imgs:    print(path)
    img = Image.open(path)
    display(img)