image_to_image.models.resfcn

  1"""
  2Module to define a simple CNN Model for image to image. 
  3
  4Classes:
  5- ResidualBlock
  6- ResFCN
  7
  8By Tobia Ippolito
  9"""
 10# ---------------------------
 11#        > Imports <
 12# ---------------------------
 13import torch
 14import torch.nn as nn
 15
 16
 17
 18# ---------------------------
 19#      > Residual FCN <
 20# ---------------------------
 21class ResidualBlock(nn.Module):
 22    """
 23    A simple residual block using fully convolutional layers.
 24
 25    The block applies two convolutional layers with an activation in between
 26    and adds the input to the output (skip connection) to preserve features.
 27    """
 28    def __init__(self, channels, kernel_size=3, padding=1):
 29        """
 30        Initializes the ResidualBlock.
 31
 32        Parameters:
 33        - channels (int): 
 34            Number of input and output channels.
 35        - kernel_size (int): 
 36            Kernel size for the convolutional layers (default=3).
 37        - padding (int): 
 38            Padding for the convolutional layers to maintain spatial dimensions (default=1).
 39        """
 40        super().__init__()
 41        self.pre_fcn = nn.Conv2d(channels, channels, kernel_size=kernel_size, padding=padding)
 42        self.activation = nn.GELU()
 43        self.post_fcn = nn.Conv2d(channels, channels, kernel_size=kernel_size, padding=padding)
 44
 45    def forward(self, x):
 46        """
 47        Forward pass through the ResidualBlock.
 48
 49        Parameters:
 50        - x (torch.tensor): 
 51            Input tensor of shape [B, C, H, W].
 52
 53        Returns:
 54        - torch.tensor: Output tensor after residual connection, same shape as input.
 55        """
 56        fcn_x = self.pre_fcn(x)
 57        fcn_x = self.activation(fcn_x)
 58        fcn_x = self.post_fcn(fcn_x)
 59        return x + fcn_x
 60
 61
 62
 63# ---------------------------
 64#        > ResFCN <
 65# ---------------------------
 66class ResFCN(nn.Module):
 67    """
 68    A simple image-to-image model using Fully Convolutional Networks (FCN) with residual connections.
 69
 70    Architecture:
 71    - Initial convolution layer to increase channel depth.
 72    - Sequence of ResidualBlocks with varying kernel sizes.
 73    - Final convolution layer to project back to output channels.
 74    - Output is clamped between 0.0 and 1.0.
 75    """
 76    def __init__(self, input_channels=1, hidden_channels=64, output_channels=1, num_blocks=64):
 77        """
 78        Initializes the ResFCN model.
 79
 80        Parameters:
 81        - input_channels (int): 
 82            Number of channels in the input image (default=1).
 83        - hidden_channels (int): 
 84            Number of channels in hidden layers / residual blocks (default=64).
 85        - output_channels (int): 
 86            Number of channels in the output image (default=1).
 87        - num_blocks (int): 
 88            Number of residual blocks to apply (default=64).
 89        """
 90        super().__init__()
 91        self.input_channels = input_channels
 92        self.output_channels = output_channels
 93
 94        self.pre_layer = nn.Conv2d(input_channels, hidden_channels, kernel_size=7, padding=3)
 95
 96        kernel_sizes = [1, 3, 5, 7, 9]
 97        residual_fcn_layers = []
 98        for i in range(num_blocks):
 99            cur_kernel_size = kernel_sizes[i%(len(kernel_sizes)-1)]
100            residual_fcn_layers += [ResidualBlock(channels=hidden_channels, 
101                                                  kernel_size=cur_kernel_size,
102                                                  padding=cur_kernel_size // 2
103                                                 )
104                                   ]
105        self.residual_fcn_layers = nn.Sequential(*residual_fcn_layers)
106
107        self.post_layer = nn.Conv2d(hidden_channels, output_channels, kernel_size=3, padding=1)
108
109    
110    def get_input_channels(self):
111        """
112        Returns the number of input channels used by the model.
113
114        Returns:
115        - int: 
116            Number of input channels expected by the model.
117        """
118        return self.input_channels
119
120    def get_output_channels(self):
121        """
122        Returns the number of output channels produced by the model.
123
124        Returns:
125        - int: 
126            Number of output channels the model generates
127        """
128        return self.output_channels
129
130    def forward(self, x):
131        """
132        Forward pass through the ResFCN model.
133
134        Parameters:
135        - x (torch.tensor): 
136            Input image tensor of shape [B, C, H, W].
137
138        Returns:
139        - torch.tensor: Output image tensor, same spatial size as input, values clamped between 0.0 and 1.0.
140        """
141        x = self.pre_layer(x)
142
143        x = self.residual_fcn_layers(x)
144
145        x = self.post_layer(x)
146
147        return torch.clamp(x, 0.0, 1.0)