src/methods/digeronimo/somefunc.py

+import numpy as np
+from scipy.signal import convolve2d
+
+from src.forward_model import CFA
+
+def bilinear_interpolation(op: CFA, z: np.ndarray) -> np.ndarray:
+    """Perform bilinear interpolation for demosaicing
+
+    Args:
+        op (CFA): CFA operator.
+        z (np.ndarray): Adjoint image.
+
+    Returns:
+        np.ndarray: Interpolated image.
+    """
+    # Bi-linear interpolation
+
+    ker_bayer_red_blue = np.array([[1, 2, 1], [2, 4, 2], [1, 2, 1]]) / 4
+    ker_bayer_green = np.array([[0, 1, 0], [1, 4, 1], [0, 1, 0]]) / 4
+
+    res = np.empty(op.input_shape)
+
+    res[:, :, 0] = convolve2d(z[:, :, 0], ker_bayer_red_blue, mode='same')
+    res[:, :, 1] = convolve2d(z[:, :, 1], ker_bayer_green, mode='same')
+    res[:, :, 2] = convolve2d(z[:, :, 2], ker_bayer_red_blue, mode='same')
+
+    return res
+
+def spectral_difference(op: CFA, z: np.ndarray, res: np.ndarray) -> np.ndarray:
+    """Perform spectral difference method for demosaicing
+
+    Args:
+        op (CFA): CFA operator.
+        z (np.ndarray): Adjoint image.
+        res (np.ndarray): Interpolated image.
+
+    Returns:
+        np.ndarray: Demosaicked image.
+    """
+
+    ker_bayer_red_blue = np.array([[1, 2, 1], [2, 4, 2], [1, 2, 1]]) / 4
+    ker_bayer_green = np.array([[0, 1, 0], [1, 4, 1], [0, 1, 0]]) / 4
+
+    # Computation of spectral differences
+
+    delta_red_green_quad =  z[:, :, 0] - np.multiply(res[:, :, 1],op.mask[:,:,0])
+    delta_red_blue_quad =  z[:, :, 0] - np.multiply(res[:, :, 2],op.mask[:,:,0])
+
+    delta_blue_green_quad =  z[:, :, 2] - np.multiply(res[:, :, 1],op.mask[:,:,2])
+    delta_blue_red_quad =  z[:, :, 2] - np.multiply(res[:, :, 0],op.mask[:,:,2])
+
+    delta_green_red_quad =  z[:, :, 1] - np.multiply(res[:, :, 0],op.mask[:,:,1])
+    delta_green_blue_quad =  z[:, :, 1] - np.multiply(res[:, :, 2],op.mask[:,:,1])
+
+    # Estimation
+    res_sd = np.empty(op.input_shape)
+
+    res_sd[:,:,0] = res[:, :, 1] + convolve2d(delta_red_green_quad,ker_bayer_red_blue,mode='same') + res[:, :, 2] + convolve2d(delta_red_blue_quad,ker_bayer_red_blue,mode='same') 
+    res_sd[:,:,2] = res[:, :, 1] + convolve2d(delta_blue_green_quad,ker_bayer_red_blue,mode='same') + res[:, :, 0] + convolve2d(delta_blue_red_quad,ker_bayer_red_blue,mode='same') 
+    res_sd[:,:,1] = res[:, :, 0] + convolve2d(delta_green_red_quad,ker_bayer_green,mode='same') + res[:, :, 2] + convolve2d(delta_green_blue_quad,ker_bayer_green,mode='same') 
+
+    return res_sd
+
+def normalization(res_sd:np.ndarray)-> np.ndarray:
+    """Perform a min-max normalization
+
+    Args:
+        res_sd (np.ndarray): Demosaicked image.
+
+    Returns:
+        np.ndarray: Normalized image.
+    """
+
+    res = (res_sd - np.min(res_sd)) / (np.max(res_sd) - np.min(res_sd))
+    
+    return res
+
+def quad_bayer_to_bayer_pattern(op: CFA):
+    """Quad to Bayer pattern conversion by swapping method
+
+    Args:
+        op (CFA): CFA operator.
+
+    Returns:
+        CFA: Bayer pettern.
+    """
+    if op.cfa == 'quad_bayer':
+        for j in range(1, op.mask.shape[1], 4):
+            op.mask[:,j], op.mask[:,j+1] = op.mask[:,j+1].copy(), op.mask[:,j].copy()
+
+        for i in range(1, op.mask.shape[0], 4):
+            op.mask[i, :], op.mask[i+1,:] = op.mask[i+1,:].copy(), op.mask[i,:].copy()
+            
+        for i in range(1, op.mask.shape[0], 4):
+            for j in range(1, op.mask.shape[1], 4):
+                op.mask[i,j], op.mask[i+1,j+1] = op.mask[i+1,j+1].copy(), op.mask[i,j].copy() 
+    return 0
+
+def quad_bayer_to_bayer_image(y:np.array):
+    """Quad to Bayer conversion of a mosaicked image by swapping method
+
+    Args:
+        y (np.array): Mosaicked image.
+    """
+        
+    for j in range(1, y.shape[1], 4):
+        y[:,j], y[:,j+1] = y[:,j+1].copy(), y[:,j].copy()
+
+    for i in range(1, y.shape[0], 4):
+        y[i, :], y[i+1,:] = y[i+1,:].copy(), y[i,:].copy()
+        
+    for i in range(1, y.shape[0], 4):
+        for j in range(1, y.shape[1], 4):
+            y[i,j], y[i+1,j+1] = y[i+1,j+1].copy(), y[i,j].copy() 
+            
+    return 0
\ No newline at end of file

src/methods/domer/forward_model.py

+"""A file containing the forward operator.
+This file should NOT be modified.
+"""
+
+
+import numpy as np
+
+from src.checks import check_cfa, check_rgb
+
+
+class CFA():
+    def __init__(self, cfa: str, input_shape: tuple) -> None:
+        """Constructor of the forward operator's class.
+
+        Args:
+            cfa (str): Name of the pattern. Either bayer or quad_bayer.
+            input_shape (tuple): Shape of the input images of the operator.
+        """
+        check_cfa(cfa)
+    
+        self.cfa = cfa
+        self.input_shape = input_shape
+        self.output_shape = input_shape[:-1]
+
+        if self.cfa == 'bayer':
+            self.mask = get_bayer_mask(input_shape)
+
+        elif self.cfa == 'quad_bayer':
+            self.mask = get_quad_bayer_mask(input_shape)
+
+
+    def direct(self, x: np.ndarray) -> np.ndarray:
+        """Applies the CFA operation to the image x.
+
+        Args:
+            x (np.ndarray): Input image.
+
+        Returns:
+            np.ndarray: Output image.
+        """
+        check_rgb(x)
+    
+        return np.sum(x * self.mask, axis=2)
+
+
+    def adjoint(self, y: np.ndarray) -> np.ndarray:
+        """Applies the adjoint of the CFA operation.
+
+        Args:
+            y (np.ndarray): Input image.
+
+        Returns:
+            np.ndarray: Output image.
+        """
+        return self.mask * y[..., np.newaxis]
+
+
+def get_bayer_mask(input_shape: tuple) -> np.ndarray:
+    """Return the mask of the Bayer CFA.
+
+    Args:
+        input_shape (tuple): Shape of the mask.
+
+    Returns:
+        np.ndarray: Mask.
+    """
+    res = np.kron(np.ones((input_shape[0], input_shape[1], 1)), [0, 1, 0])
+
+    res[::2, 1::2] = [1, 0, 0]
+    res[1::2, ::2] = [0, 0, 1]
+
+    return res
+
+
+def get_quad_bayer_mask(input_shape: tuple) -> np.ndarray:
+    """Return the mask of the quad_bayer CFA.
+
+    Args:
+        input_shape (tuple): Shape of the mask.
+
+    Returns:
+        np.ndarray: Mask.
+    """
+    res = np.kron(np.ones((input_shape[0], input_shape[1], 1)), [0, 1, 0])
+
+    res[::4, 2::4] = [1, 0, 0]
+    res[::4, 3::4] = [1, 0, 0]
+    res[1::4, 2::4] = [1, 0, 0]
+    res[1::4, 3::4] = [1, 0, 0]
+
+    res[2::4, ::4] = [0, 0, 1]
+    res[2::4, 1::4] = [0, 0, 1]
+    res[3::4, ::4] = [0, 0, 1]
+    res[3::4, 1::4] = [0, 0, 1]
+
+    return res
+
+
+####
+####
+####
+
+####      ####                ####        #############
+####      ######              ####      ##################
+####      ########            ####      ####################
+####      ##########          ####      ####        ########
+####      ############        ####      ####            ####
+####      ####  ########      ####      ####            ####
+####      ####    ########    ####      ####            ####
+####      ####      ########  ####      ####            ####
+####      ####  ##    ######  ####      ####          ######
+####      ####  ####      ##  ####      ####    ############
+####      ####  ######        ####      ####    ##########
+####      ####  ##########    ####      ####    ########
+####      ####      ########  ####      ####
+####      ####        ############      ####
+####      ####          ##########      ####
+####      ####            ########      ####
+####      ####              ######      ####
+
+# 2023
+# Authors: Mauro Dalla Mura and Matthieu Muller

src/methods/domer/image-colorization/LICENSE

+MIT License
+
+Copyright (c) 2020 priyavrat-misra
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

src/methods/domer/image-colorization/README.md

+# Image Colorization
+
+![cover](https://github.com/priyavrat-misra/image-colorization/blob/master/images/colorized/bnw_col.png?raw=true "an eye-candy")
+
+## Contents
+- [Overview](#overview)
+- [Approach](#approach)
+- [Steps](#steps)
+- [Results](#results)
+- [TL;DR](#tldr)
+- [Setup](#setup)
+- [Usage](#usage)
+- [Todo](#todo)
+
+<br>
+
+## Overview
+> This project is a Deep Convolutional Neural Network approach to solve the task of image colorization.
+> The goal is to produce a colored image given a grayscale image.<br>
+> At it's heart, it uses Convolutional Auto-Encoders to solve this task.
+> First few layers of [ResNet-18](https://arxiv.org/abs/1512.03385) model are used as the Encoder,
+> and the Decoder consists of a series of Deconvolution layers (i.e., upsample layers followed by convolutions) and residual connections.<br>
+> The model is trained on a subset of [MIT Places365](http://places2.csail.mit.edu/index.html) dataset, consisting of `41000` images of landscapes and scenes.
+
+## Approach
+> The images in the dataset are in RGB Colorspace.
+> Before loading the images, the images are converted to [LAB colorspace](https://en.wikipedia.org/wiki/CIELAB_color_space).
+> This colorspace contains exactly the same information as RGB.<br>
+> It has 3 channels, `Lightness, A and B`.
+> The lightness channel can be used as the grayscale equivalent of a colored image,
+> the rest 2 channels (A and B) contain the color information.<br>
+>
+> In a nutshell, the training process follows these steps:
+>> 1. The lightness channel is separated from the other 2 channels and used as the model's input.
+>> 2. The model predicts the A and B channels (or 'AB' for short).
+>> 3. The loss is calculated by comparing the predicted AB and the corresponding original AB of the input image.
+>
+> More about the training process can be found [here](https://github.com/priyavrat-misra/image-colorization/blob/master/train.ipynb "train.ipynb").
+
+## Steps
+> 1. [Defining a model architecture:](https://github.com/priyavrat-misra/image-colorization/blob/master/network.py "network.py")
+>    - The model follows an Auto-Encoder kind of architecture i.e., it has an `encoder` and a `decoder` part.
+>    - The encoder is used to _extract features_ of an image whereas,
+>    - the decoder is used to upsample the features. In other words, it increases the _spacial resolution_.
+>    - In here, the layers of the encoder are taken from ResNet-18 model, and the first conv layer is modified to take a single channel as input (i.e., grayscale or lightness) rather than 3 channels.
+>    - The decoder uses nearest neighbor upsampling (for increasing the spacial resolution),
+>     followed by convolutional layers (for dealing with the depth).
+>    - A more detailed visualization of the model architecture can be seen [here](https://github.com/priyavrat-misra/image-colorization/blob/master/images/architecture.png?raw=true 'after all "A picture is worth a thousand words" :)').
+> 2. [Defining a custom dataloader:](https://github.com/priyavrat-misra/image-colorization/blob/master/utils.py "utils.GrayscaleImageFolder")
+>    - when loading the images, it converts them to LAB, and returns L and AB separately.
+>    - it does few data processing tasks as well like applying tranforms and normalization.
+> 3. [Training the model:](https://github.com/priyavrat-misra/image-colorization/blob/master/train.ipynb "train.ipynb")
+>    - The model is trained for 64 epochs with [Adam Optimization](https://arxiv.org/abs/1412.6980).
+>    - For calculating the loss between the predicted AB and the original AB, Mean Squared Error is used.
+> 4. [Inference:](https://github.com/priyavrat-misra/image-colorization/blob/master/inference.ipynb "inference.ipynb")
+>    - Inference is done with unseen images and the results look promising, or should I say "natural"? :)
+
+## Results
+> ![results](https://github.com/priyavrat-misra/image-colorization/blob/master/images/results.png?raw=true)
+> _<sup>More colorized examples can be found in [here](https://github.com/priyavrat-misra/image-colorization/blob/master/images/colorized/).<sup>_
+
+## TL;DR
+> Given an image, the model can colorize it.
+
+## Setup
+- Clone and change directory:
+```bash
+git clone "https://github.com/priyavrat-misra/image-colorization.git"
+cd image-colorization/
+```
+- Dependencies:
+```bash
+pip install -r requirements.txt
+```
+
+## Usage
+```bash
+python colorize.py --img-path <path/to/image.jpg> --out-path <path/to/output.jpg> --res 360
+# or the short-way:
+python colorize.py -i <path/to/image.jpg> -o <path/to/output.jpg> -r 360
+```
+
+_Note:_
+> - As the model is trained with 224x224 images, it gives best results when `--res` is set to lower resolutions (<=480) and okay-ish when set around ~720.
+> - Setting `--res` higher than that of input image won't increase the output's quality.
+
+<br>
+
+## Todo
+- [x] define & train a model architecture
+- [x] add argparse support
+- [x] define a more residual architecture
+- [x] use pretrained resnet-18 params for the layers used in the encoder & train the model
+- [x] check how the colorization effect varies with image resolution
+- [x] separate the model from the checkpoint file to a different file
+- [x] complete README.md
+- [ ] deploy with flask
+- [ ] _after that, host it maybe?_
+
+<br>
+
+For any queries, feel free to reach me out on [LinkedIn](https://linkedin.com/in/priyavrat-misra/).
\ No newline at end of file

src/methods/domer/image-colorization/colorize.py

+import argparse
+import torch
+from torchvision.utils import save_image
+from network import ColorizeNet
+from utils import load_gray, to_rgb
+import os
+
+parser = argparse.ArgumentParser(description='colorizes an image')
+parser.add_argument('-i', '--img_path', type=str, metavar='', required=True,
+                    help='path and/or name of grayscale image to colorize')
+parser.add_argument('-r', '--res', type=int, metavar='',
+                    help='resizes the input to given resolution {default:360}')
+parser.add_argument('-o', '--out_path', type=str, metavar='', required=True,
+                    help='name to which the colorized image to be saved')
+
+device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+
+current_directory = os.getcwd()
+c_path='src/methods/domer/image-colorization/models/model.pth'
+c_path = os.path.join(current_directory, c_path)
+model = ColorizeNet()
+model.load_state_dict(
+    
+    torch.load(c_path, map_location='cpu')
+)
+
+
+def main():
+    args = parser.parse_args()
+    img_l = load_gray(args.img_path, shape=args.res)
+
+    model.eval()
+    with torch.no_grad():
+        img_ab = model(img_l)
+
+    img_rgb = to_rgb(img_l, img_ab)
+    save_image(torch.from_numpy(img_rgb.transpose(2, 0, 1)),
+               args.out_path)
+    print(f'>>> colorized image saved to "{args.out_path}"')
+
+
+if __name__ == '__main__':
+    main()

src/methods/domer/image-colorization/network.ipynb

+%% Cell type:code id: tags:
+
+``` 
+from network import ColorizeNet
+
+model = ColorizeNet()
+model
+```
+
+%% Output
+
+    ColorizeNet(
+      (encoder): Sequential(
+        (0): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
+        (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+        (2): ReLU(inplace=True)
+        (3): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
+        (4): Sequential(
+          (0): BasicBlock(
+            (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
+            (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+            (relu): ReLU(inplace=True)
+            (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
+            (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+          )
+          (1): BasicBlock(
+            (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
+            (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+            (relu): ReLU(inplace=True)
+            (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
+            (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+          )
+        )
+        (5): Sequential(
+          (0): BasicBlock(
+            (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
+            (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+            (relu): ReLU(inplace=True)
+            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
+            (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+            (downsample): Sequential(
+              (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)
+              (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+            )
+          )
+          (1): BasicBlock(
+            (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
+            (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+            (relu): ReLU(inplace=True)
+            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
+            (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+          )
+        )
+      )
+      (decoder): Sequential(
+        (0): Sequential(
+          (0): BasicBlock(
+            (conv1): Conv2d(128, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), bias=False)
+            (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+            (relu): ReLU(inplace=True)
+            (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
+            (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+            (res_conv): Conv2d(128, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
+            (res_bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+            (upsample): Upsample(scale_factor=2.0, mode=nearest)
+          )
+          (1): BasicBlock(
+            (conv1): Conv2d(64, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), bias=False)
+            (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+            (relu): ReLU(inplace=True)
+            (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
+            (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+            (activation): ReLU(inplace=True)
+          )
+        )
+        (1): Sequential(
+          (0): BasicBlock(
+            (conv1): Conv2d(64, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), bias=False)
+            (bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+            (relu): ReLU(inplace=True)
+            (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
+            (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+            (res_conv): Conv2d(64, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)
+            (res_bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+            (upsample): Upsample(scale_factor=2.0, mode=nearest)
+          )
+          (1): BasicBlock(
+            (conv1): Conv2d(32, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), bias=False)
+            (bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+            (relu): ReLU(inplace=True)
+            (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
+            (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+            (activation): ReLU(inplace=True)
+          )
+        )
+        (2): Sequential(
+          (0): BasicBlock(
+            (conv1): Conv2d(32, 2, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), bias=False)
+            (bn1): BatchNorm2d(2, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+            (relu): ReLU(inplace=True)
+            (conv2): Conv2d(2, 2, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
+            (bn2): BatchNorm2d(2, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+            (res_conv): Conv2d(32, 2, kernel_size=(1, 1), stride=(1, 1), bias=False)
+            (res_bn): BatchNorm2d(2, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+            (upsample): Upsample(scale_factor=2.0, mode=nearest)
+          )
+          (1): BasicBlock(
+            (conv1): Conv2d(2, 2, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), bias=False)
+            (bn1): BatchNorm2d(2, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+            (relu): ReLU(inplace=True)
+            (conv2): Conv2d(2, 2, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
+            (bn2): BatchNorm2d(2, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+            (activation): Sigmoid()
+          )
+        )
+      )
+    )
+
+%% Cell type:code id: tags:
+
+``` 
+import torch
+
+# to check if our model is working as expected
+model(torch.rand((2, 1, 224, 224))).shape
+```
+
+%% Output
+
+    torch.Size([2, 2, 224, 224])
+
+%% Cell type:code id: tags:
+
+``` 
+from utils import count_params
+
+count_params(model)  # no of trainable parameters
+```
+
+%% Output
+
+    1166016