nets/FAENet.py

import math
import cv2
import torch
import torch.nn as nn
import torchvision
from torch import nn
import torch

class eca_block(nn.Module):
    def __init__(self, channel, b=1, gamma=2):
        super(eca_block, self).__init__()
        kernel_size = int(abs((math.log(channel, 2) + b) / gamma))
        kernel_size = kernel_size if kernel_size % 2 else kernel_size + 1
        
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv = nn.Conv1d(1, 1, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, bias=False) 
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):

        y = self.avg_pool(x)
        y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
        y = self.sigmoid(y)
        return x * y.expand_as(x)
    
class DilatedConvNet(nn.Module):
    def __init__(self, in_channels, out_channels, dilation, padding, kernel_size):
        super(DilatedConvNet, self).__init__()
        self.dilated_conv = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding, dilation=dilation)
        self.relu = nn.ReLU(inplace=False)

    def forward(self, x):

        x = self.dilated_conv(x)
        x = self.relu(x)

        return x

class LAM(nn.Module):
    def __init__(self, ch=16):
        super().__init__()
        self.eca = eca_block(ch)
        self.conv1 = nn.Conv2d(6, 3, 3, padding=1)

    def forward(self, x):
        x = self.eca(x)
        x = self.conv1(x)
        return x

class RFEM(nn.Module):
    def __init__(
            self,
            ch_blocks=64,
            ch_mask=16,
    ):
        super().__init__()

        self.encoder = nn.Sequential(nn.Conv2d(3, 16, 3, padding=1),
                                     nn.LeakyReLU(True),
                                     nn.Conv2d(16, ch_blocks, 3, padding=1),
                                     nn.LeakyReLU(True))

        self.dconv1 = DilatedConvNet(ch_blocks,
                                  ch_blocks // 4,
                                  kernel_size=3,
                                  padding=1, dilation=1)
        self.dconv2 = DilatedConvNet(ch_blocks,
                                  ch_blocks // 4,
                                  kernel_size=3,
                                  padding=2, dilation=2)
        self.dconv3 = DilatedConvNet(ch_blocks,
                                  ch_blocks // 4,
                                  kernel_size=3,
                                  padding=3, dilation=3)
        self.dconv4 = nn.Conv2d(ch_blocks,
                                  ch_blocks // 4,
                                  kernel_size=7,
                                  padding=3)

        self.decoder = nn.Sequential(nn.Conv2d(ch_blocks, 16, 3, padding=1),
                                     nn.LeakyReLU(True),
                                     nn.Conv2d(16, 3, 3, padding=1),
                                     nn.LeakyReLU(True),
                                     )

        self.lam = LAM(ch_mask)

    def forward(self, x):
        x1 = self.encoder(x)
        x1_1 = self.dconv1(x1)
        x1_2 = self.dconv2(x1)
        x1_3 = self.dconv3(x1)
        x1_4 = self.dconv4(x1)
        x1 = torch.cat([x1_1, x1_2, x1_3, x1_4], dim=1)
        x1 = self.decoder(x1)
        out = x + x1
        out = torch.relu(out)
        mask = self.lam(torch.cat([x, out], dim=1))
        return out, mask

class ATEM(nn.Module):
    def __init__(self, in_ch=3, inter_ch=32, out_ch=3, kernel_size=3):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(in_ch, inter_ch, kernel_size, padding=kernel_size // 2),
            nn.LeakyReLU(True),
        )
        self.shift_conv = nn.Sequential(
            nn.Conv2d(in_ch, inter_ch, kernel_size, padding=kernel_size // 2))
        self.scale_conv = nn.Sequential(
            nn.Conv2d(in_ch, inter_ch, kernel_size, padding=kernel_size // 2))


        self.decoder = nn.Sequential(
            nn.Conv2d(inter_ch, out_ch, kernel_size, padding=kernel_size // 2))

    def forward(self, x, tag):
        x = self.encoder(x)
        scale = self.scale_conv(tag)
        shift = self.shift_conv(tag)
        x = x +(x * scale + shift)
        x = self.decoder(x)
        return x

class Trans_high(nn.Module):
    def __init__(self, in_ch=3, inter_ch=16, out_ch=3, kernel_size=3):
        super().__init__()
        self.atem = ATEM(in_ch, inter_ch, out_ch, kernel_size)
    def forward(self, x, tag):
        x = x + self.atem(x, tag)
        return x


class Up_tag(nn.Module):
    def __init__(self, kernel_size=1, ch=3):
        super().__init__()
        self.up = nn.Sequential(
            nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True),
            nn.Conv2d(ch,
                      ch,
                      kernel_size,
                      stride=1,
                      padding=kernel_size // 2,
                      bias=False))

    def forward(self, x):
        x = self.up(x)
        return x


class Lap_Pyramid_Conv(nn.Module):
    def __init__(self, num_high=3, kernel_size=5, channels=3):
        super().__init__()

        self.num_high = num_high
        self.kernel = self.gauss_kernel(kernel_size, channels)

    def gauss_kernel(self, kernel_size, channels):
        kernel = cv2.getGaussianKernel(kernel_size, 0).dot(
            cv2.getGaussianKernel(kernel_size, 0).T)
        kernel = torch.FloatTensor(kernel).unsqueeze(0).repeat(
            channels, 1, 1, 1)
        kernel = torch.nn.Parameter(data=kernel, requires_grad=False)
        return kernel

    def conv_gauss(self, x, kernel):
        n_channels, _, kw, kh = kernel.shape
        x = torch.nn.functional.pad(x, (kw // 2, kh // 2, kw // 2, kh // 2),
                                mode='reflect') 
        x = torch.nn.functional.conv2d(x, kernel, groups=n_channels)
        return x
    def downsample(self, x):
        return x[:, :, ::2, ::2]
    def pyramid_down(self, x):
        return self.downsample(self.conv_gauss(x, self.kernel))
    def upsample(self, x):
        up = torch.zeros((x.size(0), x.size(1), x.size(2) * 2, x.size(3) * 2),
                         device=x.device)
        up[:, :, ::2, ::2] = x * 4

        return self.conv_gauss(up, self.kernel)

    def pyramid_decom(self, img):
        self.kernel = self.kernel.to(img.device)
        current = img
        pyr = []
        for _ in range(self.num_high):
            down = self.pyramid_down(current)
            up = self.upsample(down)
            diff = current - up
            pyr.append(diff)
            current = down
        pyr.append(current)
        return pyr

    def pyramid_recons(self, pyr):
        image = pyr[0]
        for level in pyr[1:]:
            up = self.upsample(image)
            image = up + level
        return image

class FAENet(nn.Module):
    def __init__(self,
                 num_high=1,
                 ch_blocks=32,
                 up_ksize=1,
                 high_ch=32,
                 high_ksize=3,
                 ch_mask=32,
                 gauss_kernel=7):
        super().__init__()
        self.num_high = num_high
        self.lap_pyramid = Lap_Pyramid_Conv(num_high, gauss_kernel)
        self.rfem = RFEM(ch_blocks, ch_mask)

        for i in range(0, self.num_high):
            self.__setattr__('up_tag_layer_{}'.format(i),
                             Up_tag(up_ksize, ch=3))
            self.__setattr__('trans_high_layer_{}'.format(i),
                             Trans_high(3, high_ch, 3, high_ksize))

    def forward(self, x):
        pyrs = self.lap_pyramid.pyramid_decom(img=x)

        trans_pyrs = []
        trans_pyr, tag = self.rfem(pyrs[-1])
        trans_pyrs.append(trans_pyr)

        commom_tag = []
        for i in range(self.num_high):
            tag = self.__getattr__('up_tag_layer_{}'.format(i))(tag)
            commom_tag.append(tag)

        for i in range(self.num_high):
            trans_pyr = self.__getattr__('trans_high_layer_{}'.format(i))(
                pyrs[-2 - i], commom_tag[i])
            trans_pyrs.append(trans_pyr)

        out = self.lap_pyramid.pyramid_recons(trans_pyrs)

        return out
    

faenet = FAENet()
params = faenet.parameters()
num_params = sum(p.numel() for p in params)
print("FAENet parameters: {:.2f}K   ".format(num_params/ 1024) + "{:.2f} MB".format(num_params/ (1024 * 1024)))