{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"},{"sourceId":125589,"sourceType":"datasetVersion","datasetId":50512},{"sourceId":1462296,"sourceType":"datasetVersion","datasetId":857191},{"sourceId":11449504,"sourceType":"datasetVersion","datasetId":7088240},{"sourceId":106675505,"sourceType":"kernelVersion"},{"sourceId":106676815,"sourceType":"kernelVersion"}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"class Config:\n    external_data_size = 10000\n    train_size = 8000\n    image_size_1 = 256\n    image_size_2 = 256\n    batch_size = 32\n    LeakyReLU_slope = 0.2\n    dropout = 0.5\n    kernel_size = 4\n    stride = 2\n    padding = 1\n    gen_lr = 2e-4\n    disc_lr = 2e-4\n    beta1 = 0.5\n    beta2 = 0.999\n    lambda_l1 = 100\n    gan_mode = 'vanilla'\n    layers_to_cut = -2\n    epochs = 15\n    pretrain_lr = 1e-4","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# !pip install --upgrade torch torchvision","metadata":{"_kg_hide-output":true,"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nfrom pathlib import Path\nimport glob\nimport time\nimport numpy as np\nimport matplotlib.pyplot as plt\n%matplotlib inline\nfrom tqdm import tqdm\n\nimport PIL\nfrom PIL import Image\nfrom skimage.color import rgb2lab, lab2rgb\n\nimport torch\nfrom torch import nn, optim\nfrom torchvision import transforms\nfrom torchvision.models.resnet import resnet18\nfrom torchvision.models.vgg import vgg19\nfrom torch.utils.data import Dataset, DataLoader\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","metadata":{"_kg_hide-output":true,"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## FastAI imports + Loading in the Data\n\nThe data is a small sample of the COCO dataset for object detection and is loaded in through the external data functionality of fastai. The actual paper uses the entire ImageNet dataset while we only use 8000 training images.","metadata":{}},{"cell_type":"code","source":"!pip install fastai==2.4\nfrom fastai.vision.learner import create_body\nfrom fastai.vision.models.unet import DynamicUnet\nfrom fastai.data.external import untar_data, URLs\n\npath = untar_data(URLs.COCO_SAMPLE)\npath = str(path) + \"/train_sample\"\npaths = glob.glob(path + \"/*.jpg\")","metadata":{"_kg_hide-output":true,"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Train test split\nnp.random.seed(42)\npaths_subset = np.random.choice(paths, Config.external_data_size, replace=False) # choosing 10000 images randomly\nrandom_idxs = np.random.permutation(Config.external_data_size) \ntrain_idxs = random_idxs[:Config.train_size] # choosing the first 8000 as training set\nval_idxs = random_idxs[Config.train_size:] # choosing last 2000 as validation set\ntrain_paths = paths_subset[train_idxs]\nval_paths = paths_subset[val_idxs]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Taking a look at the training dataset\n_, axes = plt.subplots(4, 4, figsize=(10, 10))\nfor ax, img_path in zip(axes.flatten(), train_paths):\n    ax.imshow(Image.open(img_path))\n    ax.axis(\"off\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ImageDataset(Dataset):\n    ''' \n    \n    Class that deals with the Data Loading and preprocessing steps such as image resizing, data augmentation (horizontal flip) \n    and conversion of RGB image to LAB color space with standardization.\n    \n    '''\n    \n    def __init__(self,paths,train = True):\n        if train == True:\n            self.transforms = transforms.Compose([transforms.Resize((Config.image_size_1,Config.image_size_2)),\n                                                 transforms.RandomHorizontalFlip()]) # Basic Data Augmentation\n        elif train == False:\n            self.transforms = transforms.Compose([transforms.Resize((Config.image_size_1,Config.image_size_2))])\n            \n        self.train = train\n        self.paths = paths\n        \n    def __len__(self):\n        return len(self.paths)\n    \n    def __getitem__(self,idx):\n        img = Image.open(self.paths[idx]).convert(\"RGB\")\n        img = self.transforms(img)\n        img = np.array(img)\n        lab = rgb2lab(img).astype(\"float32\")\n        lab = transforms.ToTensor()(lab)\n        L = lab[[0],...]/50 - 1 # Standardizing L space\n        ab = lab[[1,2],...]/128 # Standardizing ab space\n        \n        return {'L': L,\n                'ab': ab}","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"batch_size = Config.batch_size\n\n#train dataset\ntrain_data = ImageDataset(paths = train_paths, train=True)\n# validation dataset\nvalid_data = ImageDataset(paths = val_paths, train=False)\n# train data loader\ntrain_loader = DataLoader(train_data, batch_size=batch_size,shuffle=True,pin_memory = True)\n# validation data loader\nvalid_loader = DataLoader(valid_data, batch_size=batch_size,shuffle=False,pin_memory = True)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\nclass AverageMeter:\n    '''\n    \n    Computes and stores the average and current value for the various losses involved.\n    \n    '''\n    \n    def __init__(self):\n        self.reset()\n        \n    def reset(self):\n        self.count, self.avg, self.sum = [0.] * 3\n    \n    def update(self, val, count=1):\n        self.count += count\n        self.sum += count * val\n        self.avg = self.sum / self.count\n\ndef create_loss_meters():\n    disc_loss_gen = AverageMeter()\n    disc_loss_real = AverageMeter()\n    disc_loss = AverageMeter()\n    loss_G_GAN = AverageMeter()\n    loss_G_L1 = AverageMeter()\n    loss_G = AverageMeter()\n    \n    return {'disc_loss_gen': disc_loss_gen,\n            'disc_loss_real': disc_loss_real,\n            'disc_loss': disc_loss,\n            'loss_G_GAN': loss_G_GAN,\n            'loss_G_L1': loss_G_L1,\n            'loss_G': loss_G}\n\ndef update_losses(model, loss_meter_dict, count):\n    for loss_name, loss_meter in loss_meter_dict.items():\n        loss = getattr(model, loss_name)\n        loss_meter.update(loss.item(), count=count)\n        \n\n\ndef lab_to_rgb(L, ab):\n    \"\"\"\n    \n    Takes a batch of images and converts from LAB space to RGB\n    \n    \"\"\"\n    \n    L = (L + 1.) * 50\n    ab = ab * 128\n    Lab = torch.cat([L, ab], dim=1).permute(0, 2, 3, 1).cpu().numpy()\n    rgb_imgs = []\n    for img in Lab:\n        img_rgb = lab2rgb(img)\n        rgb_imgs.append(img_rgb)\n    return np.stack(rgb_imgs, axis=0)\n    \ndef visualize(model, data, save=True):\n    '''\n    \n    Function which helps visualize and see the difference between the input image (grayscale), generated image \n    and the Real image.\n    \n    '''\n    \n    model.generator.eval()\n    with torch.no_grad():\n        model.prepare_input(data)\n        model.forward()\n#     model.generator.train()\n    fake_color = model.gen_output.detach()\n    real_color = model.ab\n    L = model.L\n    fake_imgs = lab_to_rgb(L, fake_color)\n    real_imgs = lab_to_rgb(L, real_color)\n    fig = plt.figure(figsize=(15, 8))\n    for i in range(5):\n        ax = plt.subplot(3, 5, i + 1)\n        ax.imshow(L[i][0].cpu(), cmap='gray')\n        ax.axis(\"off\")\n        ax = plt.subplot(3, 5, i + 1 + 5)\n        ax.imshow(fake_imgs[i])\n        ax.axis(\"off\")\n        ax = plt.subplot(3, 5, i + 1 + 10)\n        ax.imshow(real_imgs[i])\n        ax.axis(\"off\")\n    plt.show()\n    if save:\n        fig.savefig(f\"colorization_{time.time()}.png\")\n        \ndef log_results(loss_meter_dict):\n    for loss_name, loss_meter in loss_meter_dict.items():\n        print(f\"{loss_name}: {loss_meter.avg:.5f}\")\n        \n\nclass GANLoss(nn.Module):\n    '''\n    \n    Class that defines the GAN loss.\n    \n    '''\n    \n    def __init__(self, gan_mode='vanilla', real_label=1.0, fake_label=0.0):\n        super().__init__()\n        self.register_buffer('real_label', torch.tensor(real_label))\n        self.register_buffer('fake_label', torch.tensor(fake_label))\n        if gan_mode == 'vanilla':\n            self.loss = nn.BCEWithLogitsLoss()\n        elif gan_mode == 'lsgan':\n            self.loss = nn.MSELoss()\n    \n    def get_labels(self, preds, target_is_real):\n        if target_is_real:\n            labels = self.real_label\n        else:\n            labels = self.fake_label\n        return labels.expand_as(preds)\n    \n    def __call__(self, preds, target_is_real):\n        labels = self.get_labels(preds, target_is_real)\n        loss = self.loss(preds, labels)\n        return loss","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Code help for model architecture and training loop: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix\n\nclass UnetBlock(nn.Module):\n    '''\n    \n    Class that helps define the various different types of blocks used in the Generator architecture\n    \n    '''\n    \n    def __init__(self, nf, ni, submodule=None, input_channels=None, dropout=False,innermost=False, outermost=False):\n        super().__init__()\n        self.outermost = outermost\n        if input_channels is None: \n            input_channels = nf\n        downconv = nn.Conv2d(input_channels, ni, kernel_size=Config.kernel_size, stride=Config.stride, padding=Config.padding, bias=False)\n        downrelu = nn.LeakyReLU(Config.LeakyReLU_slope, True)\n        downnorm = nn.BatchNorm2d(ni)\n        uprelu = nn.ReLU(True)\n        upnorm = nn.BatchNorm2d(nf)\n        \n        if outermost:\n            upconv = nn.ConvTranspose2d(ni * 2, nf,kernel_size=Config.kernel_size, stride=Config.stride, padding=Config.padding)\n            down = [downconv]\n            up = [uprelu, upconv, nn.Tanh()]\n            model = down + [submodule] + up\n        elif innermost:\n            upconv = nn.ConvTranspose2d(ni, nf, kernel_size=Config.kernel_size, stride=Config.stride, padding=Config.padding, bias=False)\n            down = [downrelu, downconv]\n            up = [uprelu, upconv, upnorm]\n            model = down + up\n        else:\n            upconv = nn.ConvTranspose2d(ni * 2, nf, kernel_size=Config.kernel_size, stride=Config.stride, padding=Config.padding, bias=False)\n            down = [downrelu, downconv, downnorm]\n            up = [uprelu, upconv, upnorm]\n            if dropout: up += [nn.Dropout(Config.dropout)]\n            model = down + [submodule] + up\n        self.model = nn.Sequential(*model)\n    \n    def forward(self, x):\n        if self.outermost:\n            return self.model(x)\n        else:\n            return torch.cat([x, self.model(x)], 1)\n\nclass Unet(nn.Module):\n    '''\n    \n    The Generator Model Class\n    \n    '''\n    \n    def __init__(self, input_channels=1, output_channels=2, n_down=8, num_filters=64):\n        super().__init__()\n        unet_block = UnetBlock(num_filters * 8, num_filters * 8, innermost=True)\n        for _ in range(n_down - 5):\n            unet_block = UnetBlock(num_filters * 8, num_filters * 8, submodule=unet_block, dropout=True)\n        out_filters = num_filters * 8\n        for _ in range(3):\n            unet_block = UnetBlock(out_filters // 2, out_filters, submodule=unet_block)\n            out_filters //= 2\n        self.model = UnetBlock(output_channels, out_filters, input_channels=input_channels, submodule=unet_block, outermost=True)\n    \n    def forward(self, x):\n        return self.model(x)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    '''\n    \n    The Discriminator Model Class\n    \n    '''\n    \n    def __init__(self,input_channels,num_filters=64, n_down=3):\n        super().__init__()\n        model = [self.get_layers(input_channels,num_filters,norm = False)]\n        model += [self.get_layers(num_filters * 2 ** i, num_filters * 2 ** (i + 1), stride=1 if i == (n_down-1) else 2) for i in range(n_down)]\n        model += [self.get_layers(num_filters * 2 ** n_down, 1, stride=1, norm=False, activation=False)] \n        self.model = nn.Sequential(*model)\n    \n    def get_layers(self, ni, nf, kernel_size=Config.kernel_size, stride=Config.stride, padding=Config.padding, norm=True, activation=True): \n        layers = [nn.Conv2d(ni, nf, kernel_size, stride, padding, bias=not norm)]          \n        if norm: \n            layers += [nn.BatchNorm2d(nf)]\n        if activation:\n            layers += [nn.LeakyReLU(Config.LeakyReLU_slope, True)]\n        return nn.Sequential(*layers)\n    \n    def forward(self, x):\n        return self.model(x)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Normal Weight Initialization with mean = 0, std = 0.02\n\ndef init_weights(net, init='norm', gain=0.02):\n    \n    def init_func(m):\n        classname = m.__class__.__name__\n        if hasattr(m, 'weight') and 'Conv' in classname:\n            if init == 'norm':\n                nn.init.normal_(m.weight.data, mean=0.0, std=gain)\n            elif init == 'xavier':\n                nn.init.xavier_normal_(m.weight.data, gain=gain)\n            elif init == 'kaiming':\n                nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')\n            \n            if hasattr(m, 'bias') and m.bias is not None:\n                nn.init.constant_(m.bias.data, 0.0)\n        elif 'BatchNorm2d' in classname:\n            nn.init.normal_(m.weight.data, 1., gain)\n            nn.init.constant_(m.bias.data, 0.)\n            \n    net.apply(init_func)\n    print(f\"model initialized with {init} initialization\")\n    return net\n\ndef init_model(model, device):\n    model = model.to(device)\n    model = init_weights(model)\n    return model","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def build_backbone_unet(input_channels=1, output_channels=2, size=Config.image_size_1):\n    '''\n    Function that helps build the generator backbone using another pretrained model trained on imagenet for image classification\n    '''\n    \n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    resnet_model = resnet18(pretrained=True)\n    body = create_body(resnet_model, pretrained=True, n_in=input_channels, cut=Config.layers_to_cut)\n    generator = DynamicUnet(body, output_channels, (size, size)).to(device)\n    return generator\n\ndef pretrain_generator(generator, train_dl, opt, criterion, epochs):\n    '''\n    \n    Function to pretrain the generator in a supervised manner on the colorization task using L1 Loss.\n    \n    '''\n    \n    for e in range(epochs):\n        loss_meter = AverageMeter()\n        for data in tqdm(train_dl):\n            L, ab = data['L'].to(device), data['ab'].to(device)\n            preds = generator(L)\n            loss = criterion(preds, ab)\n            opt.zero_grad()\n            loss.backward()\n            opt.step()\n            \n            loss_meter.update(loss.item(), L.size(0))\n            \n        print(f\"Epoch {e + 1}/{epochs}\")\n        print(f\"L1 Loss: {loss_meter.avg:.5f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class MainModel(nn.Module):\n    '''\n    \n    The Class where it all comes together.\n    \n    '''\n    \n    def __init__(self,generator = None,gen_lr =Config.gen_lr,disc_lr = Config.disc_lr, beta1 = Config.beta1,beta2 = Config.beta2, lambda_l1 = Config.lambda_l1):\n        super().__init__()\n        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n        self.lambda_l1 = lambda_l1\n        \n        if generator is None:\n            self.generator = init_model(Unet(input_channels=1, output_channels=2, n_down=8, num_filters=64),self.device)\n        else:\n            self.generator = generator.to(self.device)\n        self.discriminator = init_model(Discriminator(input_channels = 3,num_filters = 64,n_down = 3),self.device)\n        self.GANloss = GANLoss(gan_mode=Config.gan_mode).to(self.device)\n        self.L1loss = nn.L1Loss()\n        self.gen_optim = optim.Adam(self.generator.parameters(), lr=gen_lr, betas=(beta1, beta2))\n        self.disc_optim= optim.Adam(self.discriminator.parameters(), lr=disc_lr, betas=(beta1, beta2))\n        \n    def requires_grad(self,model,requires_grad = True):\n        for p in model.parameters():\n            p.requires_grad = requires_grad\n            \n    def prepare_input(self,data):\n        self.L = data['L'].to(self.device)\n        self.ab = data['ab'].to(self.device)\n        \n    def forward(self):\n        self.gen_output = self.generator(self.L)\n        \n    def disc_backward(self):\n        gen_image = torch.cat([self.L, self.gen_output], dim=1)\n        gen_image_preds = self.discriminator(gen_image.detach())\n        self.disc_loss_gen = self.GANloss(gen_image_preds, False)\n        real_image = torch.cat([self.L, self.ab], dim=1)\n        real_preds = self.discriminator(real_image)\n        self.disc_loss_real = self.GANloss(real_preds, True)\n        self.disc_loss = (self.disc_loss_gen + self.disc_loss_real) * 0.5\n        self.disc_loss.backward()\n    \n    def gen_backward(self):\n        gen_image = torch.cat([self.L, self.gen_output], dim=1)\n        gen_image_preds = self.discriminator(gen_image)\n        self.loss_G_GAN = self.GANloss(gen_image_preds, True)\n        self.loss_G_L1 = self.L1loss(self.gen_output, self.ab) * self.lambda_l1\n        self.loss_G = self.loss_G_GAN + self.loss_G_L1\n        self.loss_G.backward()\n    \n    def optimize(self):\n        self.forward()\n        self.discriminator.train()\n        self.requires_grad(self.discriminator, True)\n        self.disc_optim.zero_grad()\n        self.disc_backward()\n        self.disc_optim.step()\n        \n        self.generator.train()\n        self.requires_grad(self.discriminator, False) \n        self.gen_optim.zero_grad()\n        self.gen_backward()\n        self.gen_optim.step()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def train_model(model, train_loader, epochs, display=100):\n    data = next(iter(train_loader)) \n    for epoch in range(epochs):\n        loss_meter_dict = create_loss_meters() \n        i=0\n        for data in tqdm(train_loader):\n            model.prepare_input(data) \n            model.optimize()\n            update_losses(model, loss_meter_dict, count=data['L'].size(0)) \n            i+=1\n            if i % display == 0:\n                print(f\"\\nEpoch {epoch+1}/{epochs}\")\n                print(f\"Iteration {i}/{len(train_loader)}\")\n                log_results(loss_meter_dict) \n                visualize(model, data, save=False) ","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"generator = build_backbone_unet(input_channels=1, output_channels=2, size=Config.image_size_1)\nopt = optim.Adam(generator.parameters(), lr=Config.pretrain_lr)\nloss = nn.L1Loss()        \npretrain_generator(generator, train_loader, opt, loss, Config.epochs)\ntorch.save(generator.state_dict(), \"res18-unet.pt\")\ngenerator.load_state_dict(torch.load(\"res18-unet.pt\", map_location=device))\nmodel = MainModel(generator=generator)\ntrain_model(model, train_loader, Config.epochs)\ntorch.save(model.state_dict(), \"main-model.pt\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport torch\nfrom skimage.metrics import peak_signal_noise_ratio\n\ndef calculate_psnr(real_images, generated_images, data_range=1.0):\n    \"\"\"\n    Calculate PSNR between real and generated images.\n    \n    Args:\n        real_images: Tensor of real images, shape [B, C, H, W], values in [0, 1]\n        generated_images: Tensor of generated images, shape [B, C, H, W], values in [0, 1]\n        data_range: The data range of the input image (default: 1.0)\n        \n    Returns:\n        Mean PSNR value across the batch\n    \"\"\"\n    # Convert tensors to numpy arrays if they're not already\n    if isinstance(real_images, torch.Tensor):\n        real_images = real_images.detach().cpu().numpy()\n    if isinstance(generated_images, torch.Tensor):\n        generated_images = generated_images.detach().cpu().numpy()\n    \n    # Ensure images are in the right shape [B, C, H, W]\n    assert real_images.shape == generated_images.shape, \"Real and generated image shapes must match\"\n    \n    # Calculate PSNR for each image pair\n    psnr_values = []\n    for i in range(real_images.shape[0]):\n        # Move channel to last dim for skimage: [C, H, W] -> [H, W, C]\n        real_img = np.transpose(real_images[i], (1, 2, 0))\n        gen_img = np.transpose(generated_images[i], (1, 2, 0))\n        \n        # Calculate PSNR\n        psnr = peak_signal_noise_ratio(real_img, gen_img, data_range=data_range)\n        psnr_values.append(psnr)\n    \n    return np.mean(psnr_values)\n\n# Example usage with the pix2pix model:\ndef evaluate_psnr(model, data_loader, device):\n    \"\"\"\n    Evaluate PSNR on a dataset using the trained model\n    \"\"\"\n    model.generator.eval()\n    psnr_list = []\n    \n    with torch.no_grad():\n        for data in data_loader:\n            # Get L and ab channels\n            L = data['L'].to(device)\n            real_ab = data['ab'].to(device)\n            \n            # Generate colorization\n            fake_ab = model.generator(L)\n            \n            # Convert from LAB to RGB\n            L_unorm = (L + 1.) * 50.  # Unstandardize L channel\n            real_ab_unorm = real_ab * 128.  # Unstandardize ab channels\n            fake_ab_unorm = fake_ab * 128.  # Unstandardize ab channels\n            \n            # Concatenate channels to get LAB images\n            real_lab = torch.cat([L_unorm, real_ab_unorm], dim=1)\n            fake_lab = torch.cat([L_unorm, fake_ab_unorm], dim=1)\n            \n            # Convert LAB to RGB\n            real_rgb = []\n            fake_rgb = []\n            for i in range(real_lab.size(0)):\n                real_rgb_img = lab_to_rgb(real_lab[i:i+1])\n                fake_rgb_img = lab_to_rgb(fake_lab[i:i+1])\n                real_rgb.append(real_rgb_img)\n                fake_rgb.append(fake_rgb_img)\n            \n            real_rgb = np.concatenate(real_rgb, axis=0)\n            fake_rgb = np.concatenate(fake_rgb, axis=0)\n            \n            # Calculate PSNR\n            batch_psnr = calculate_psnr(real_rgb, fake_rgb)\n            psnr_list.append(batch_psnr)\n    \n    return np.mean(psnr_list)\n\n# This function is taken from the provided notebook to convert LAB to RGB\ndef lab_to_rgb(lab_image):\n    lab_image = lab_image.detach().cpu().numpy()\n    lab_image = np.transpose(lab_image[0], (1, 2, 0))  # [C, H, W] -> [H, W, C]\n    \n    # Suppress the specific warning about negative Z values\n    import warnings\n    with warnings.catch_warnings():\n        warnings.filterwarnings(\"ignore\", message=\"Conversion from CIE-LAB, via XYZ to sRGB color space resulted in.*negative Z values\")\n        rgb_image = lab2rgb(lab_image)\n    \n    return rgb_image[np.newaxis, ...]  # Add batch dimension back","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport torch\nfrom skimage.metrics import structural_similarity as ssim\n\ndef calculate_ssim(real_images, generated_images, data_range=1.0, win_size=3):\n    \"\"\"\n    Calculate SSIM between real and generated images using a small window size.\n    \n    Args:\n        real_images: Numpy array of real images, shape [B, C, H, W], values in [0, 1]\n        generated_images: Numpy array of generated images, shape [B, C, H, W], values in [0, 1]\n        data_range: The data range of the input image (default: 1.0)\n        win_size: Window size for SSIM calculation, must be odd (default: 3)\n        \n    Returns:\n        Mean SSIM value across the batch\n    \"\"\"\n    # Convert tensors to numpy arrays if they're not already\n    if isinstance(real_images, torch.Tensor):\n        real_images = real_images.detach().cpu().numpy()\n    if isinstance(generated_images, torch.Tensor):\n        generated_images = generated_images.detach().cpu().numpy()\n    \n    # Ensure images are in the right shape [B, C, H, W]\n    assert real_images.shape == generated_images.shape, \"Real and generated image shapes must match\"\n    \n    # Calculate SSIM for each image pair\n    ssim_values = []\n    for i in range(real_images.shape[0]):\n        # Move channel to last dim for skimage: [C, H, W] -> [H, W, C]\n        real_img = np.transpose(real_images[i], (1, 2, 0))\n        gen_img = np.transpose(generated_images[i], (1, 2, 0))\n        \n        # Get current image dimensions\n        h, w, c = real_img.shape\n        \n        # Skip if image is too small for even the smallest window\n        if h < 3 or w < 3:\n            continue\n            \n        # Use the smallest possible window size (3) or user-specified size if smaller than image\n        current_win_size = min(win_size, h - 1 if h % 2 == 0 else h, w - 1 if w % 2 == 0 else w)\n        # Ensure window size is odd\n        if current_win_size % 2 == 0:\n            current_win_size -= 1\n            \n        try:\n            # Calculate SSIM with safe window size\n            ssim_val = ssim(real_img, gen_img, \n                           data_range=data_range,\n                           channel_axis=-1 if c > 1 else None,  # Newer scikit-image uses channel_axis\n                           win_size=current_win_size)\n            ssim_values.append(ssim_val)\n        except Exception as e:\n            print(f\"Error calculating SSIM for image {i}: {e}\")\n            print(f\"Image shape: {real_img.shape}, win_size: {current_win_size}\")\n            # Continue with next image rather than failing\n            continue\n    \n    # If no valid calculations, return 0\n    if len(ssim_values) == 0:\n        return 0.0\n        \n    return np.mean(ssim_values)\n\n# Example usage with the pix2pix model:\ndef evaluate_ssim(model, data_loader, device, win_size=3):\n    \"\"\"\n    Evaluate SSIM on a dataset using the trained model\n    \n    Args:\n        model: The pix2pix model\n        data_loader: DataLoader for evaluation data\n        device: Device to run the model on\n        win_size: Window size for SSIM calculation (default: 3)\n        \n    Returns:\n        Mean SSIM value across the dataset\n    \"\"\"\n    model.generator.eval()\n    ssim_list = []\n    \n    with torch.no_grad():\n        for data in data_loader:\n            # Get L and ab channels\n            L = data['L'].to(device)\n            real_ab = data['ab'].to(device)\n            \n            # Generate colorization\n            fake_ab = model.generator(L)\n            \n            # Convert from LAB to RGB\n            L_unorm = (L + 1.) * 50.  # Unstandardize L channel\n            real_ab_unorm = real_ab * 128.  # Unstandardize ab channels\n            fake_ab_unorm = fake_ab * 128.  # Unstandardize ab channels\n            \n            # Concatenate channels to get LAB images\n            real_lab = torch.cat([L_unorm, real_ab_unorm], dim=1)\n            fake_lab = torch.cat([L_unorm, fake_ab_unorm], dim=1)\n            \n            # Convert LAB to RGB\n            real_rgb = []\n            fake_rgb = []\n            \n            # Suppress the specific warning about negative Z values\n            import warnings\n            with warnings.catch_warnings():\n                warnings.filterwarnings(\"ignore\", message=\"Conversion from CIE-LAB, via XYZ to sRGB color space resulted in.*negative Z values\")\n                \n                for i in range(real_lab.size(0)):\n                    real_rgb_img = lab_to_rgb(real_lab[i:i+1])\n                    fake_rgb_img = lab_to_rgb(fake_lab[i:i+1])\n                    real_rgb.append(real_rgb_img)\n                    fake_rgb.append(fake_rgb_img)\n            \n            real_rgb = np.concatenate(real_rgb, axis=0)\n            fake_rgb = np.concatenate(fake_rgb, axis=0)\n            \n            # Calculate SSIM with the specified window size\n            batch_ssim = calculate_ssim(real_rgb, fake_rgb, win_size=win_size)\n            ssim_list.append(batch_ssim)\n    \n    return np.mean(ssim_list)\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport numpy as np\nfrom scipy import linalg\nfrom torchvision import models, transforms\nfrom torch.utils.data import DataLoader, Dataset\n\nclass InceptionV3Features(nn.Module):\n    \"\"\"Inception v3 model for computing FID features\"\"\"\n    def __init__(self):\n        super().__init__()\n        # Load pretrained Inception v3 model\n        inception = models.inception_v3(pretrained=True)\n        # We only need features before the final classification layer\n        self.blocks = nn.Sequential(\n            inception.Conv2d_1a_3x3, inception.Conv2d_2a_3x3, inception.Conv2d_2b_3x3,\n            nn.MaxPool2d(kernel_size=3, stride=2),\n            inception.Conv2d_3b_1x1, inception.Conv2d_4a_3x3,\n            nn.MaxPool2d(kernel_size=3, stride=2),\n            inception.Mixed_5b, inception.Mixed_5c, inception.Mixed_5d,\n            inception.Mixed_6a, inception.Mixed_6b, inception.Mixed_6c, inception.Mixed_6d,\n            inception.Mixed_6e, inception.Mixed_7a, inception.Mixed_7b, inception.Mixed_7c,\n            nn.AdaptiveAvgPool2d(output_size=(1, 1))\n        )\n        # Set model to evaluation mode\n        self.eval()\n        # Freeze parameters\n        for param in self.parameters():\n            param.requires_grad = False\n            \n    def forward(self, x):\n        # Expected input: batches of RGB images with values in [0, 1]\n        # Preprocess to match Inception v3 expectations\n        # Resize to 299x299 and normalize\n        x = nn.functional.interpolate(x, size=(299, 299), mode='bilinear', align_corners=False)\n        x = x * 2 - 1  # Scale from [0, 1] to [-1, 1]\n        # Extract features\n        features = self.blocks(x)\n        features = features.squeeze(3).squeeze(2)  # Shape: [batch_size, 2048]\n        return features\n\ndef calculate_fid(real_features, fake_features):\n    \"\"\"\n    Calculate FID score between real and fake feature distributions.\n    \n    Args:\n        real_features: Features from real images, shape [N, d]\n        fake_features: Features from fake images, shape [N, d]\n        \n    Returns:\n        FID score (lower is better)\n    \"\"\"\n    # Convert to numpy\n    if isinstance(real_features, torch.Tensor):\n        real_features = real_features.cpu().numpy()\n    if isinstance(fake_features, torch.Tensor):\n        fake_features = fake_features.cpu().numpy()\n    \n    # Calculate mean and covariance\n    mu_real = np.mean(real_features, axis=0)\n    mu_fake = np.mean(fake_features, axis=0)\n    sigma_real = np.cov(real_features, rowvar=False)\n    sigma_fake = np.cov(fake_features, rowvar=False)\n    \n    # Calculate FID\n    mu_diff = mu_real - mu_fake\n    covmean, _ = linalg.sqrtm(sigma_real @ sigma_fake, disp=False)\n    if np.iscomplexobj(covmean):\n        covmean = covmean.real\n    \n    fid = mu_diff @ mu_diff + np.trace(sigma_real + sigma_fake - 2 * covmean)\n    return fid\n\ndef evaluate_fid(model, data_loader, device, num_samples=1000):\n    \"\"\"\n    Evaluate FID score on a dataset using the trained model\n    \n    Args:\n        model: The pix2pix model\n        data_loader: DataLoader for evaluation data\n        device: Device to run the model on\n        num_samples: Number of samples to use for FID calculation\n        \n    Returns:\n        FID score (lower is better)\n    \"\"\"\n    model.generator.eval()\n    inception_model = InceptionV3Features().to(device)\n    \n    # Collect real and fake features\n    real_features = []\n    fake_features = []\n    sample_count = 0\n    \n    with torch.no_grad():\n        for data in data_loader:\n            if sample_count >= num_samples:\n                break\n                \n            # Get L and ab channels\n            L = data['L'].to(device)\n            real_ab = data['ab'].to(device)\n            batch_size = L.size(0)\n            \n            # Generate colorization\n            fake_ab = model.generator(L)\n            \n            # Convert to RGB images for Inception\n            # Convert from LAB to RGB\n            L_unorm = (L + 1.) * 50.  # Unstandardize L channel\n            real_ab_unorm = real_ab * 128.  # Unstandardize ab channels\n            fake_ab_unorm = fake_ab * 128.  # Unstandardize ab channels\n            \n            # Concatenate channels to get LAB images\n            real_lab = torch.cat([L_unorm, real_ab_unorm], dim=1)\n            fake_lab = torch.cat([L_unorm, fake_ab_unorm], dim=1)\n            \n            # Convert LAB to RGB tensors\n            real_rgb_batch = []\n            fake_rgb_batch = []\n            \n            # Suppress the specific warning about negative Z values\n            import warnings\n            with warnings.catch_warnings():\n                warnings.filterwarnings(\"ignore\", message=\"Conversion from CIE-LAB, via XYZ to sRGB color space resulted in.*negative Z values\")\n                \n                for i in range(batch_size):\n                    lab_img_real = real_lab[i].permute(1, 2, 0).cpu().numpy()\n                    lab_img_fake = fake_lab[i].permute(1, 2, 0).cpu().numpy()\n                    \n                    # Convert from LAB to RGB using lab2rgb\n                    rgb_img_real = lab2rgb(lab_img_real)\n                    rgb_img_fake = lab2rgb(lab_img_fake)\n                    \n                    # Convert back to tensor [H, W, C] -> [C, H, W]\n                    rgb_tensor_real = torch.from_numpy(rgb_img_real).permute(2, 0, 1).float().to(device)\n                    rgb_tensor_fake = torch.from_numpy(rgb_img_fake).permute(2, 0, 1).float().to(device)\n                    \n                    real_rgb_batch.append(rgb_tensor_real)\n                    fake_rgb_batch.append(rgb_tensor_fake)\n            \n            # Stack batch\n            real_rgb_batch = torch.stack(real_rgb_batch)\n            fake_rgb_batch = torch.stack(fake_rgb_batch)\n            \n            # Extract features\n            batch_real_features = inception_model(real_rgb_batch)\n            batch_fake_features = inception_model(fake_rgb_batch)\n            \n            real_features.append(batch_real_features)\n            fake_features.append(batch_fake_features)\n            \n            sample_count += batch_size\n    \n    # Concatenate all features\n    real_features = torch.cat(real_features, dim=0)\n    fake_features = torch.cat(fake_features, dim=0)\n    \n    # Calculate FID\n    fid_score = calculate_fid(real_features[:num_samples], fake_features[:num_samples])\n    \n    return fid_score\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Evaluate the model\npsnr_score = evaluate_psnr(model, valid_loader, device)\nssim_score = evaluate_ssim(model, valid_loader, device)\nfid_score = evaluate_fid(model, valid_loader, device)\n\nprint(f\"PSNR: {psnr_score:.4f} dB\")  # Higher is better\nprint(f\"SSIM: {ssim_score:.4f}\")     # Higher is better (0-1 range)\nprint(f\"FID: {fid_score:.4f}\")       # Lower is better","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom PIL import Image\nfrom skimage.color import rgb2lab, lab2rgb\nfrom torchvision import transforms\n\n# Define path to the grayscale image\ngrayscale_image_path = \"/kaggle/input/ffffffff/Grayscaleimage42785.jpg\"\n\n# Define paths to your saved models\ngenerator_path = \"res18-unet.pt\"  # Path to the pretrained generator\nmodel_path = \"main-model.pt\"     # Path to the full pix2pix model\n\ndef preprocess_grayscale_image(image_path, image_size=256):\n    \"\"\"\n    Load and preprocess a grayscale image for colorization.\n    Returns the L channel in the LAB color space.\n    \"\"\"\n    # Load image\n    img = Image.open(image_path).convert(\"RGB\")\n    \n    # Resize image\n    img = transforms.Resize((image_size, image_size))(img)\n    \n    # Convert to numpy array\n    img = np.array(img)\n    \n    # Convert to LAB\n    lab_img = rgb2lab(img).astype(\"float32\")\n    \n    # Extract L channel\n    L = lab_img[:,:,0]\n    \n    # Normalize L channel to [-1, 1]\n    L = L / 50.0 - 1.0\n    \n    # Add channel and batch dimensions\n    L = torch.from_numpy(L).unsqueeze(0).unsqueeze(0)\n    \n    return L, img\n\ndef colorize_image(model, L, device=\"cuda\"):\n    \"\"\"\n    Generate colorized image from L channel using the pix2pix model.\n    \"\"\"\n    # Ensure model is in eval mode\n    model.generator.eval()\n    \n    with torch.no_grad():\n        # Move to device\n        L = L.to(device)\n        \n        # Generate ab channels\n        ab = model.generator(L)\n        \n        # Prepare for conversion back to RGB\n        L_unorm = (L + 1.) * 50.  # Unstandardize L\n        ab_unorm = ab * 128.      # Unstandardize ab\n        \n        # Concatenate channels\n        lab = torch.cat([L_unorm, ab_unorm], dim=1)\n        \n        # Convert to numpy\n        lab_np = lab[0].cpu().numpy().transpose(1, 2, 0)  # [C, H, W] -> [H, W, C]\n        \n        # Suppress warning about negative Z values\n        import warnings\n        with warnings.catch_warnings():\n            warnings.filterwarnings(\"ignore\", message=\"Conversion from CIE-LAB, via XYZ to sRGB color space resulted in.*negative Z values\")\n            # Convert LAB to RGB\n            rgb_img = lab2rgb(lab_np)\n        \n    return rgb_img\n\n# Load the trained model\ndef load_model(model_path, device=\"cuda\"):\n    \"\"\"Load the trained pix2pix model.\"\"\"\n    from __main__ import MainModel, build_backbone_unet  # Access classes from notebook\n    \n    # Two options for loading the model\n    try:\n        print(\"Attempting to load model with resnet18 backbone...\")\n        # Option 1: Create generator with resnet18 backbone first\n        generator = build_backbone_unet(input_channels=1, output_channels=2, size=256)\n        model = MainModel(generator=generator)\n        \n        # Load the model state\n        state_dict = torch.load(model_path, map_location=device)\n        model.load_state_dict(state_dict)\n        print(\"Successfully loaded model with resnet18 backbone.\")\n        \n    except Exception as e1:\n        print(f\"First loading attempt failed: {e1}\")\n        try:\n            print(\"Attempting to load generator directly...\")\n            # Option 2: Try loading just the generator\n            generator_model = build_backbone_unet(input_channels=1, output_channels=2, size=256)\n            generator_model.load_state_dict(torch.load(model_path, map_location=device))\n            \n            # Create a simplified model with just the loaded generator\n            class SimpleColorizer:\n                def __init__(self, generator):\n                    self.generator = generator\n                    self.generator.eval()\n            \n            model = SimpleColorizer(generator_model)\n            print(\"Successfully loaded generator only.\")\n            \n        except Exception as e2:\n            print(f\"Second loading attempt failed: {e2}\")\n            print(\"Falling back to direct model loading...\")\n            # Option 3: Try direct loading\n            model = MainModel()\n            model.load_state_dict(torch.load(model_path, map_location=device), strict=False)\n            print(\"Loaded model with missing keys (some parameters may be uninitialized).\")\n    \n    return model\n\n# Put it all together\ndef predict_colorization(image_path, model_path):\n    # Determine device\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    print(f\"Using device: {device}\")\n    \n    # Load and preprocess the grayscale image\n    L, original_img = preprocess_grayscale_image(image_path)\n    \n    # Load the model\n    model = load_model(model_path, device)\n    \n    # Generate colorized image\n    colorized_img = colorize_image(model, L, device)\n    \n    # Display results\n    plt.figure(figsize=(12, 6))\n    \n    plt.subplot(1, 3, 1)\n    plt.imshow(original_img)\n    plt.title(\"Original Image\")\n    plt.axis(\"off\")\n    \n    plt.subplot(1, 3, 2)\n    plt.imshow(original_img, cmap='gray')\n    plt.title(\"Grayscale Image\")\n    plt.axis(\"off\")\n    \n    plt.subplot(1, 3, 3)\n    plt.imshow(colorized_img)\n    plt.title(\"Colorized Image\")\n    plt.axis(\"off\")\n    \n    plt.tight_layout()\n    plt.show()\n    \n    return colorized_img\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\ngrayscale_image_path = \"/kaggle/input/ffffffff/Grayscaleimage42785.jpg\"\nprint(\"Trying to load and predict with the generator model...\")\ncolorized_result = predict_colorization(grayscale_image_path, generator_path)\n\n\n# Save the colorized image if needed\nplt.imsave(\"colorized_output.jpg\", colorized_result)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom PIL import Image\nfrom diffusers import StableDiffusionUpscalePipeline\n\ndef upscale_image(input_image_path, prompt=None, output_path=\"upscaled_image.png\"):\n\n    # Use a default prompt if none provided\n    if prompt is None:\n        prompt = \"A clear, detailed, high-resolution photograph with natural colors\"\n    \n    print(\"Loading Stable Diffusion X4 Upscaler model...\")\n    # Load the upscaler pipeline\n    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n    pipe = StableDiffusionUpscalePipeline.from_pretrained(\n        \"stabilityai/stable-diffusion-x4-upscaler\", \n        torch_dtype=torch.float16 if device == \"cuda\" else torch.float32\n    )\n    pipe = pipe.to(device)\n    \n    # Load and prepare the image\n    print(\"Loading input image...\")\n    low_res_img = Image.open(input_image_path).convert(\"RGB\")\n    \n    # Get original dimensions for display\n    original_width, original_height = low_res_img.size\n    print(f\"Original image dimensions: {original_width}x{original_height}\")\n    \n    # The model expects images around 128-256px, so resize if needed\n    # while maintaining the aspect ratio\n    if max(original_width, original_height) > 256:\n        print(\"Resizing input image to fit model requirements...\")\n        if original_width > original_height:\n            new_width = 256\n            new_height = int(original_height * (256 / original_width))\n        else:\n            new_height = 256\n            new_width = int(original_width * (256 / original_height))\n        \n        low_res_img = low_res_img.resize((new_width, new_height), Image.LANCZOS)\n        print(f\"Resized to: {new_width}x{new_height}\")\n    \n    # Upscale the image using the pipeline\n    print(\"Upscaling image with Stable Diffusion (this may take a while)...\")\n    upscaled_image = pipe(prompt=prompt, image=low_res_img).images[0]\n    \n    # Get upscaled dimensions\n    upscaled_width, upscaled_height = upscaled_image.size\n    print(f\"Upscaled image dimensions: {upscaled_width}x{upscaled_height}\")\n    \n    # Save the upscaled image\n    upscaled_image.save(output_path)\n    print(f\"Upscaled image saved to: {output_path}\")\n    \n    # Display comparison\n    plt.figure(figsize=(12, 6))\n    \n    plt.subplot(1, 2, 1)\n    plt.imshow(np.array(low_res_img))\n    plt.title(f\"Original ({original_width}x{original_height})\")\n    plt.axis(\"off\")\n    \n    plt.subplot(1, 2, 2)\n    plt.imshow(np.array(upscaled_image))\n    plt.title(f\"Upscaled ({upscaled_width}x{upscaled_height})\")\n    plt.axis(\"off\")\n    \n    plt.tight_layout()\n    plt.show()\n    \n    return upscaled_image\n\n# Path to the colorized output from the previous step\ncolorized_output_path = \"colorized_output.jpg\"\n\n# Upscale with a custom prompt that matches the content of your image\n# Modify this prompt based on what your image actually contains for better results\ncustom_prompt = \"A high resolution photograph with natural colors, detailed textures, and sharp features ,8k\"\n\n# Run the upscaling\nupscaled_result = upscale_image(\n    input_image_path=colorized_output_path,\n    prompt=custom_prompt,\n    output_path=\"upscaled_colorized_output.png\"\n)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}