model.py

import torch
import torch.nn as nn 
from torchvision.models import resnet18 , ResNet18_Weights



class ColorizationNet(nn.Module):
    def __init__(self):
        super().__init__()

        # Convert grayscale (1 channel) to RGB (3 channels)
        self.gray_to_rgb = nn.Conv2d(
            in_channels=1,
            out_channels=3,
            kernel_size=3,
            padding=1
        )

        # Pretrained ResNet18 Encoder
        backbone = resnet18(weights=ResNet18_Weights.DEFAULT)
        self.encoder = nn.Sequential(*list(backbone.children())[:-2])

        # Freeze encoder weights
        for param in self.encoder.parameters():
            param.requires_grad = False

        # Decoder
        self.decoder = nn.Sequential(

            nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),

            nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),

            nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),

            nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),

            nn.ConvTranspose2d(32, 2, kernel_size=2, stride=2),
            nn.Tanh()
        )

    def forward(self,x):
        x = self.gray_to_rgb(x)
        with torch.no_grad():
            x = self.encoder(x)

        x= self.decoder(x)
        return x