论文笔记之【Iterative Filter Adaptive Network for Single Image Defocus Deblurring】


代码网址:https://github.com/codeslake/IFAN
一、摘要

提出新型迭代自适应滤波器网络IFAN,它主要功能是处理空间的变化和大面积离焦像素。为了将预测可分离滤波器应用于特征,我们也提出 IAC层(Iterative Adaptive Convolution)迭代地将可分离滤波器应用于特征。为了精确的使图像模糊,IFAN需要预测离焦滤波器的准确信息,其中包含离焦模糊的大小和形状。

二、数据集

训练集:使用DPDD数据集的双像素图像训练网络(双像素图像提供了一堆带有微小基线的立体图像,其差异与离焦模糊程度成正比)。

文中可用训练集:DPDD dataset、Pixel dual-pixel test set、CUHK blur detection dataset

三、实验流程:

(一)IAC(Iterative Adaptive Convolution)网络简介

IAC网络是受FAC网络(补充内容)的启发研究出来的。

IAC网络结构图如下:

输入:特征图e和空间过滤器F,他们拥有一样的尺寸大小。在F中的每个空间位置是Nc(2k+1)维,其中N代表有N组滤波器{F1,F2,....,FN}。第n个滤波器集合Fn有两个1维滤波器f1n和f2n,他们的尺寸是k×1,和1×k,还有一个偏差向量bn,f1n和f2n和bn均有c个通道。IAC层将F中每个位置

代表应用了第n个离散过滤器和偏置得到的第n个特征图的。IAC中的可以用如下式计算。其中LReLU(leaky rectified linear unit)为激活函数。

IAC网络代码

import torch
import torch.nn.functional as Func

def IAC(feat_in, F, N, c, k, is_act_last = True):
    Fs = torch.split(F[:, :N * (c * k * 2), :, :], c * k * 2, dim = 1)
    F_bs = torch.split(F[:, N * (c * k * 2):, :, :], c, dim = 1)
    for i in range(N):
        F1, F2= torch.split(Fs[i], c * k, dim = 1)
        f = SAC(feat_in=feat_in if i == 0 else f, kernel1=F1, kernel2=F2, ksize=k) ## image
        f = f + F_bs[i]

        if i < (N - 1):
            f = Func.leaky_relu(f, 0.1, inplace=True)
        elif is_act_last:
            f = Func.leaky_relu(f, 0.1, inplace=True)

    return f

def SAC(feat_in, kernel1, kernel2, ksize):
    channels = feat_in.size(1)
    N, kernels, H, W = kernel1.size()
    pad = (ksize - 1) // 2

    feat_in = Func.pad(feat_in, (0, 0, pad, pad), mode="replicate")
    feat_in = feat_in.unfold(2, ksize, 1)
    feat_in = feat_in.permute(0, 2, 3, 1, 4).view(N, H, W, channels, -1)

    kernel1 = kernel1.permute(0, 2, 3, 1).view(N, H, W, channels, ksize)
    feat_in = torch.sum(torch.mul(feat_in, kernel1), -1)
    feat_in = feat_in.permute(0, 3, 1, 2)

    feat_in = Func.pad(feat_in, (pad, pad, 0, 0), mode="replicate")
    feat_in = feat_in.unfold(3, ksize, 1)
    feat_in = feat_in.permute(0, 2, 3, 1, 4).view(N, H, W, channels, -1)
    kernel2 = kernel2.permute(0, 2, 3, 1).view(N, H, W, channels, ksize)
    feat_in = torch.sum(torch.mul(feat_in, kernel1), -1)
    feat_out = feat_in.permute(0, 3, 1, 2)

    return feat_out

  (二)IFAN去模糊网络

IFAN的网络架构图如下图所示:

 输入:单幅离焦模糊图像IB∈RH×W×3。其中H和W分别代表图像的高和宽。

网络组成:这个网路由简单的特征提取,重构和中间的IFAN网络构成的编码-解码器。 特征提取到的特征图eB∈Rh×w×ce。其中h=H/8,w=W/8。IFAN网络通过预测空间离散模糊滤波器和将eB应用于IAC网络中得到去除离焦模糊的特征图eBS。再将eBS传送到重构图中,的到全聚焦图像IBS。

IFAN代码:

import torch
import torch.nn as nn
import torch.nn.functional as Func
import collections

from models.utils import DPD
from models.nn_common import conv, upconv, resnet_block
from models.IAC import IAC

class Network(nn.Module):
    def __init__(self, config):
        super(Network, self).__init__()
        try:
            self.rank = torch.distributed.get_rank()
        except Exception as ex:
            self.rank = 0

        ks = config.ks # kernel size for convolution
        self.Fs = config.Fs # kernel size for IAC
        res_num = config.res_num

        ch1 = config.ch
        ch2 = ch1 * 2
        ch3 = ch1 * 4
        ch4 = ch1 * 4
        self.ch4 = ch4

        # weight init for filter predictor
        self.wiF = config.wiF

        ###################################
        # Feature Extractor - Reconstructor
        ###################################
        # feature extractor
        self.conv1_1 = conv(3, ch1, kernel_size=ks, stride=1)
        self.conv1_2 = conv(ch1, ch1, kernel_size=ks, stride=1)
        self.conv1_3 = conv(ch1, ch1, kernel_size=ks, stride=1)

        self.conv2_1 = conv(ch1, ch2, kernel_size=ks, stride=2)
        self.conv2_2 = conv(ch2, ch2, kernel_size=ks, stride=1)
        self.conv2_3 = conv(ch2, ch2, kernel_size=ks, stride=1)

        self.conv3_1 = conv(ch2, ch3, kernel_size=ks, stride=2)
        self.conv3_2 = conv(ch3, ch3, kernel_size=ks, stride=1)
        self.conv3_3 = conv(ch3, ch3, kernel_size=ks, stride=1)

        self.conv4_1 = conv(ch3, ch4, kernel_size=ks, stride=2)
        self.conv4_2 = conv(ch4, ch4, kernel_size=ks, stride=1)
        self.conv4_3 = conv(ch4, ch4, kernel_size=ks, stride=1)

        self.conv4_4 = nn.Sequential(
            conv(2 * ch4, ch4, kernel_size=ks),
            resnet_block(ch4, kernel_size=ks, res_num=res_num),
            resnet_block(ch4, kernel_size=ks, res_num=res_num),
            conv(ch4, ch4, kernel_size=ks))

        # reconstructor
        self.conv_res = nn.Sequential(
            conv(ch4, ch4, kernel_size=ks),
            resnet_block(ch4, kernel_size=ks, res_num=3),
            conv(ch4, ch4, kernel_size=ks))

        self.upconv3_u = upconv(ch4, ch3)
        self.upconv3_1 = resnet_block(ch3, kernel_size=ks, res_num=1)
        self.upconv3_2 = resnet_block(ch3, kernel_size=ks, res_num=1)

        self.upconv2_u = upconv(ch3, ch2)
        self.upconv2_1 = resnet_block(ch2, kernel_size=ks, res_num=1)
        self.upconv2_2 = resnet_block(ch2, kernel_size=ks, res_num=1)

        self.upconv1_u = upconv(ch2, ch1)
        self.upconv1_1 = resnet_block(ch1, kernel_size=ks, res_num=1)
        self.upconv1_2 = resnet_block(ch1, kernel_size=ks, res_num=1)

        self.out_res = conv(ch1, 3, kernel_size=ks)
        ###################################

        ###################################
        # IFAN
        ###################################
        # filter encoder
        self.kconv1_1 = conv(3, ch1, kernel_size=ks, stride=1)
        self.kconv1_2 = conv(ch1, ch1, kernel_size=ks, stride=1)
        self.kconv1_3 = conv(ch1, ch1, kernel_size=ks, stride=1)

        self.kconv2_1 = conv(ch1, ch2, kernel_size=ks, stride=2)
        self.kconv2_2 = conv(ch2, ch2, kernel_size=ks, stride=1)
        self.kconv2_3 = conv(ch2, ch2, kernel_size=ks, stride=1)

        self.kconv3_1 = conv(ch2, ch3, kernel_size=ks, stride=2)
        self.kconv3_2 = conv(ch3, ch3, kernel_size=ks, stride=1)
        self.kconv3_3 = conv(ch3, ch3, kernel_size=ks, stride=1)

        self.kconv4_1 = conv(ch3, ch4, kernel_size=ks, stride=2)
        self.kconv4_2 = conv(ch4, ch4, kernel_size=ks, stride=1)
        self.kconv4_3 = conv(ch4, ch4, kernel_size=ks, stride=1)

        # disparity map estimator
        self.DME = nn.Sequential(
            conv(ch4, ch4, kernel_size=ks),
            resnet_block(ch4, kernel_size=ks, res_num=res_num),
            resnet_block(ch4, kernel_size=ks, res_num=res_num),
            conv(ch4, 1, kernel_size=3, act = None))

        # filter predictor
        self.conv_DME = conv(1, ch4, kernel_size=3)
        self.N = config.N
        self.kernel_dim = self.N * (ch4 * self.Fs * 2) + self.N * ch4
        self.F = nn.Sequential(
            conv(ch4, ch4, kernel_size=ks),
            resnet_block(ch4, kernel_size=ks, res_num=res_num),
            resnet_block(ch4, kernel_size=ks, res_num=res_num),
            conv(ch4, self.kernel_dim, kernel_size=1, act = None))

    def weights_init_F(self, m):
        if isinstance(m, torch.nn.Conv2d) or isinstance(m, torch.nn.ConvTranspose2d):
            torch.nn.init.xavier_uniform_(m.weight, gain = self.wiF)
            if m.bias is not None:
                torch.nn.init.constant_(m.bias, 0)

    def init_F(self):
        self.F.apply(self.weights_init_F)

##########################################################################
    def forward(self, C, R=None, L=None, is_train=False):
        # feature extractor
        f1 = self.conv1_3(self.conv1_2(self.conv1_1(C)))
        f2 = self.conv2_3(self.conv2_2(self.conv2_1(f1)))
        f3 = self.conv3_3(self.conv3_2(self.conv3_1(f2)))
        f_C = self.conv4_3(self.conv4_2(self.conv4_1(f3)))

        # filter encoder
        f = self.kconv1_3(self.kconv1_2(self.kconv1_1(C)))
        f = self.kconv2_3(self.kconv2_2(self.kconv2_1(f)))
        f = self.kconv3_3(self.kconv3_2(self.kconv3_1(f)))
        f = self.kconv4_3(self.kconv4_2(self.kconv4_1(f)))

        # disparity map estimator
        DM = self.DME(f)

        # filter predictor
        f_DM = self.conv_DME(DM)
        f = self.conv4_4(torch.cat([f, f_DM], 1))
        F = self.F(f)

        # IAC
        f = IAC(f_C, F, self.N, self.ch4, self.Fs)

        # reconstructor
        f = self.conv_res(f)

        f = self.upconv3_u(f) + f3
        f = self.upconv3_2(self.upconv3_1(f))

        f = self.upconv2_u(f) + f2
        f = self.upconv2_2(self.upconv2_1(f))

        f = self.upconv1_u(f) + f1
        f = self.upconv1_2(self.upconv1_1(f))

        out = self.out_res(f) + C

        # results
        outs = collections.OrderedDict()

        if is_train is False:
            outs['result'] = torch.clip(out, 0, 1.0)
        else:
            outs['result'] = out
            # F
            outs['Filter'] = F

            # DME
            f = self.kconv1_3(self.kconv1_2(self.kconv1_1(R)))
            f = self.kconv2_3(self.kconv2_2(self.kconv2_1(f)))
            f = self.kconv3_3(self.kconv3_2(self.kconv3_1(f)))
            f = self.kconv4_3(self.kconv4_2(self.kconv4_1(f)))
            DM = self.DME(f)
            f_R_warped = DPD(Func.interpolate(R, scale_factor=1/8, mode='area'), DM, padding_mode = 'zeros')
            outs['f_R_w'] = f_R_warped
            outs['f_L'] = Func.interpolate(L, scale_factor=1/8, mode='area')

        return outs 


四、效果图

 五、最终结论

提出一种基于迭代滤波自适应网络(IFAN)的单图像离焦去模糊框架。IFAN通过预测每像素可分离去模糊滤波器,灵活的处理空间变化的离焦模糊。为有效的处理较大的模糊区域,提出迭代自适应卷积(IAC)迭代地在特征上应用可分离滤波器。此外,IFAN通过学习视差图估计和重模糊来更准确地估计单张图像的离焦模糊。

六、补充内容:

(一)FAC网络

FAC网络架构图如下图所示:

输入:特征图e(h×w)和空间过滤器F(h×w)。其中特征图e有c个通道,ck2维的空间过滤器。

过程:在使用过程中,每一个空间点(x,y)在FAC层中生成c个卷积滤波器,然后该层以通道方式将滤波器应用于以(x,y)为中心的e中的特征,从而产生一个输出特征图e∈Rh×w×c。