sscp.py

import torch
import torch.nn as nn
import torch.nn.functional as F


# Squeeze and Excitation block
class SELayer(nn.Module):
    def __init__(self, num_channels, reduction_ratio=8):
        '''
            num_channels: The number of input channels
            reduction_ratio: The reduction ratio 'r' from the paper
        '''
        super(SELayer, self).__init__()
        self.num_channels_reduced = num_channels // reduction_ratio
        self.reduction_ratio = reduction_ratio
        self.fc1 = nn.Linear(num_channels, num_channels, bias=True)
        self.fc2 = nn.Linear(num_channels, num_channels, bias=True)
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()

    def forward(self, input_tensor):

        batch_size, num_channels, H, W = input_tensor.size()

        squeeze_tensor = input_tensor.view(batch_size, num_channels, -1).mean(dim=2)

        # channel excitation
        fc_out_1 = self.relu(self.fc1(squeeze_tensor))
        fc_out_2 = self.sigmoid(self.fc2(fc_out_1))

        a, b = squeeze_tensor.size()
        output_tensor = torch.mul(input_tensor, fc_out_2.view(a, b, 1, 1))
        return output_tensor


# SSPCAB implementation
class SSPCAB(nn.Module):
    def __init__(self, channels, kernel_dim=1, dilation=1, reduction_ratio=8):
        '''
            channels: The number of filter at the output (usually the same with the number of filter from the input)
            kernel_dim: The dimension of the sub-kernels ' k' ' from the paper
            dilation: The dilation dimension 'd' from the paper
            reduction_ratio: The reduction ratio for the SE block ('r' from the paper)
        '''
        super(SSPCAB, self).__init__()
        self.pad = kernel_dim + dilation
        self.border_input = kernel_dim + 2*dilation + 1

        self.mse_loss = nn.MSELoss()

        self.relu = nn.ReLU()
        self.se = SELayer(channels, reduction_ratio=reduction_ratio)

        self.conv1 = nn.Conv2d(in_channels=channels,
                               out_channels=channels,
                               kernel_size=kernel_dim)
        self.conv2 = nn.Conv2d(in_channels=channels,
                               out_channels=channels,
                               kernel_size=kernel_dim)
        self.conv3 = nn.Conv2d(in_channels=channels,
                               out_channels=channels,
                               kernel_size=kernel_dim)
        self.conv4 = nn.Conv2d(in_channels=channels,
                               out_channels=channels,
                               kernel_size=kernel_dim)

    def forward(self, input):
        x =  F.pad(input, (self.pad, self.pad, self.pad, self.pad), "constant", 0)
        x1 = self.conv1(x[:, :, :-self.border_input, :-self.border_input])
        x2 = self.conv2(x[:, :, self.border_input:, :-self.border_input])
        x3 = self.conv3(x[:, :, :-self.border_input, self.border_input:])
        x4 = self.conv4(x[:, :, self.border_input:, self.border_input:])
        x = self.relu(x1 + x2 + x3 + x4)
        output = self.se(x)
        #loss=self.forward_loss(input,output)
        return output
    
    
    def forward_loss(self,input,predicted):
            
        cost_sspcab = self.mse_loss(input,predicted)
        return cost_sspcab

class ResConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ResConvBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out += residual
        return out

class ConvolutionalModel(nn.Module):
    def __init__(self, num_blocks=2, num_channels=64):
        super(ConvolutionalModel, self).__init__()
        self.initial_conv = nn.Conv2d(3, num_channels, kernel_size=3, padding=1)
        self.bn = nn.BatchNorm2d(num_channels)
        self.relu = nn.ReLU()
        self.mse_loss = nn.MSELoss()

        layers = []
        for _ in range(num_blocks):
            layers.append(ResConvBlock(num_channels, num_channels))
        self.res_blocks = nn.Sequential(*layers)

        self.penumlimate_conv=SSPCAB(num_channels,kernel_dim=3,dilation=5, reduction_ratio=8)
        self.last_conv = nn.Conv2d(num_channels, 3, kernel_size=3, padding=1)

    def forward(self, x):
        initial_out = self.initial_conv(x)
        out = self.relu(initial_out)
        out = self.res_blocks(out)
        out=self.penumlimate_conv(out)
        out=self.bn(out)
        out += initial_out  # Skip connection from input to output feature map
        out = self.last_conv(out)
        loss=self.forward_loss(x,out)
        return out,loss
    
    def forward_loss(self,input,predicted):     
        cost_sspcab = self.mse_loss(input,predicted)
        return cost_sspcab