image-colorization-using-gans.ipynb

{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.7.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat_minor":4,"nbformat":4,"cells":[{"source":"<a href=\"https://www.kaggle.com/code/prabhanjanjadhav/image-colorization-using-gans?scriptVersionId=114252738\" target=\"_blank\"><img align=\"left\" alt=\"Kaggle\" title=\"Open in Kaggle\" src=\"https://kaggle.com/static/images/open-in-kaggle.svg\"></a>","metadata":{},"cell_type":"markdown","outputs":[],"execution_count":0},{"cell_type":"markdown","source":"<span><h1 style = \"font-family: garamond; font-size: 40px; font-style: normal; letter-spcaing: 3px; background-color: #f6f5f5; color :#fe346e; border-radius: 100px 100px; text-align:center\">Image Colorization using GANs 🧨</h1></span> \n\nIn this notebook you'll learn to implement GANs, unets, and realize how pretraining is used in improve the model performance.\nThis notebook tackles the task of image colorization using conditional GANs. The implementation is inspired from pix2pix paper.\n# <span><h1 style = \"font-family: garamond; font-size: 40px; font-style: normal; letter-spcaing: 3px; background-color: #f6f5f5; color :#fe346e; border-radius: 100px 100px; text-align:center\">Intro: Generative Adversarial Networks 🍭<a id=\"Intro\"></a></h1></span>\nThis module contains brief introduction to GANs. Feel free to skip it if you're familiar with it.\n\nA GAN model comprises of 2 sub-models:\n\n* The <b>generator</b> learns to generate plausible data. The generated instances become negative training examples for the discriminator.\n* The <b>discriminator</b> learns to distinguish the generator's fake data from real data. The discriminator penalizes the generator for producing implausible results.\n\nWe simultaneously train two models: a generative model G that captures the data distribution, and a discriminative model D that estimates\nthe probability that a sample came from the training data rather than G. The training procedure for G is to maximize the probability of D making a mistake. This framework corresponds to a minimax two-player game. In the space of arbitrary functions G and D, a unique solution exists, with G recovering the training data distribution and D equal to ${1 \\over 2}$ everywhere. ","metadata":{}},{"cell_type":"markdown","source":"\n<img src=\"https://media.springernature.com/full/springer-static/image/art%3A10.1038%2Fs41524-020-00352-0/MediaObjects/41524_2020_352_Fig1_HTML.png\" width=\"800\"></img>\n","metadata":{}},{"cell_type":"markdown","source":"# <span><h1 style = \"font-family: garamond; font-size: 40px; font-style: normal; letter-spcaing: 3px; background-color: #f6f5f5; color :#fe346e; border-radius: 100px 100px; text-align:center\">Install Required Libraries 📦 <a id=\"Installing required libraries\"></a></h1></span>","metadata":{}},{"cell_type":"code","source":"# !pip install --upgrade torch torchvision","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2022-12-07T09:20:16.326503Z","iopub.execute_input":"2022-12-07T09:20:16.326893Z","iopub.status.idle":"2022-12-07T09:20:16.332463Z","shell.execute_reply.started":"2022-12-07T09:20:16.326853Z","shell.execute_reply":"2022-12-07T09:20:16.331172Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# <span><h1 style = \"font-family: garamond; font-size: 40px; font-style: normal; letter-spcaing: 3px; background-color: #f6f5f5; color :#fe346e; border-radius: 100px 100px; text-align:center\">Import Required Libraries 🧺<a id=\"Importing libraries\"></a></h1></span>","metadata":{}},{"cell_type":"code","source":"import os\nimport glob\nimport time\n\n# For data manipulation\nimport numpy as np\nfrom PIL import Image\nimport cv2 as cv\nfrom pathlib import Path\n\n# Pytorch imports\nimport torch\nfrom torch import nn, optim\nfrom torchvision import transforms\nfrom torchvision.utils import make_grid\nfrom torch.utils.data import Dataset, DataLoader\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n\n# Utils\nfrom tqdm.notebook import tqdm\nimport matplotlib.pyplot as plt\nfrom skimage.color import rgb2lab, lab2rgb","metadata":{"id":"zPRfy4p4bKDU","execution":{"iopub.status.busy":"2022-12-07T09:17:50.997042Z","iopub.execute_input":"2022-12-07T09:17:50.997585Z","iopub.status.idle":"2022-12-07T09:17:51.007097Z","shell.execute_reply.started":"2022-12-07T09:17:50.997547Z","shell.execute_reply":"2022-12-07T09:17:51.006142Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install fastai==2.4\nfrom fastai.vision.learner import create_body\nfrom torchvision.models.resnet import resnet18\nfrom fastai.vision.models.unet import DynamicUnet","metadata":{"id":"Snladw4hbKDX","outputId":"e6084aef-50a9-4778-ac26-24474f377657","scrolled":true,"execution":{"iopub.status.busy":"2022-12-07T09:18:01.86535Z","iopub.execute_input":"2022-12-07T09:18:01.866401Z","iopub.status.idle":"2022-12-07T09:19:27.282669Z","shell.execute_reply.started":"2022-12-07T09:18:01.866365Z","shell.execute_reply":"2022-12-07T09:19:27.281517Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# <span><h1 style = \"font-family: garamond; font-size: 40px; font-style: normal; letter-spcaing: 3px; background-color: #f6f5f5; color :#fe346e; border-radius: 100px 100px; text-align:center\">Making Dataset and Dataloaders🥡<a id=\"Datasets and Dataloaders\"></a></h1></span>\nWe'll be using a subset of <b>COCO dataset</b>: 8000 training images and 2000 validation images.","metadata":{}},{"cell_type":"code","source":"# setting seed\nnp.random.seed(123)\n\npaths = glob.glob('/kaggle/input/coco-2017-dataset/coco2017/train2017' + \"/*.jpg\") # Grabbing all the image file paths\npaths_subset = np.random.choice(paths, 10_000, replace=False) # choosing 10000 paths randomly\nrand_idxs = np.random.permutation(10_000) # generate a numpy array of numbers from 0 to 9999 in any random order\ntrain_idxs = rand_idxs[:8000]\nval_idxs = rand_idxs[8000:]\ntrain_paths = paths_subset[train_idxs]\nval_paths = paths_subset[val_idxs]\n\nprint(len(train_paths), len(val_paths))","metadata":{"id":"vH_Ftv_yfEmX","outputId":"af59fa2f-aee3-42b5-9e13-67ebacdcb002","execution":{"iopub.status.busy":"2022-12-07T09:19:37.905925Z","iopub.execute_input":"2022-12-07T09:19:37.906341Z","iopub.status.idle":"2022-12-07T09:19:39.301906Z","shell.execute_reply.started":"2022-12-07T09:19:37.906304Z","shell.execute_reply":"2022-12-07T09:19:39.300865Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Visualizing the 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":{"id":"1HPHpZH9bKDY","outputId":"3b57c1ac-f5a2-458e-eb82-09e1fb6a7bff","execution":{"iopub.status.busy":"2022-12-07T09:19:42.466652Z","iopub.execute_input":"2022-12-07T09:19:42.467951Z","iopub.status.idle":"2022-12-07T09:19:44.267963Z","shell.execute_reply.started":"2022-12-07T09:19:42.467901Z","shell.execute_reply":"2022-12-07T09:19:44.266529Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"SIZE = 256\nclass ColorizationDataset(Dataset):\n    def __init__(self, paths, split='train'):\n        if split == 'train':\n            self.transforms = transforms.Compose([                  \n                transforms.Resize((SIZE, SIZE),  Image.BICUBIC),\n                transforms.RandomHorizontalFlip(), \n            ])\n        elif split == 'val':\n            self.transforms = transforms.Resize((SIZE, SIZE),  Image.BICUBIC)\n        \n        self.split = split\n        self.size = SIZE\n        self.paths = paths\n    \n    def __getitem__(self, idx): # function for getting L and ab tensors of an image idx\n        img = Image.open(self.paths[idx]).convert(\"RGB\") # converting images to RGB to tackle any grayscale image if present.\n        img = self.transforms(img)\n        img = np.array(img) \n        img_lab = rgb2lab(img).astype(\"float32\") # Converting RGB to L*a*b; (h,w,c)\n        img_lab = transforms.ToTensor()(img_lab) # Converting Lab img to tensor; also converts it to (c, h, w)\n        L = img_lab[[0], ...] / 50. - 1.\n        ab = img_lab[[1, 2], ...] / 110.\n        \n        return {'L': L, 'ab': ab}\n    \n    def __len__(self):\n        return len(self.paths)\n\ndef make_dataloaders(batch_size=16, n_workers=2, pin_memory=True, **kwargs): # a handy function to make our dataloaders\n                                                              \n    dataset = ColorizationDataset(**kwargs)\n    dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=n_workers,     \n                            pin_memory=pin_memory)\n    return dataloader","metadata":{"id":"awzDs1s2bKDZ","execution":{"iopub.status.busy":"2022-12-07T09:19:56.587361Z","iopub.execute_input":"2022-12-07T09:19:56.588393Z","iopub.status.idle":"2022-12-07T09:19:56.60973Z","shell.execute_reply.started":"2022-12-07T09:19:56.588353Z","shell.execute_reply":"2022-12-07T09:19:56.608526Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dl = make_dataloaders(paths=train_paths, split='train')\nval_dl = make_dataloaders(paths=val_paths, split='val')\nprint(len(train_dl)) # 8000 / dataloader_batch_size(=16) \nprint(len(val_dl))  # 2000 / 16\ndata = next(iter(train_dl))\nLs, abs_ = data['L'], data['ab']\nprint(Ls.shape, abs_.shape)","metadata":{"id":"3GcqDwnpbKDa","outputId":"60082adb-6ff6-4b41-8934-9929e9ca3d9b","execution":{"iopub.status.busy":"2022-12-07T09:19:56.865814Z","iopub.execute_input":"2022-12-07T09:19:56.866788Z","iopub.status.idle":"2022-12-07T09:20:01.356825Z","shell.execute_reply.started":"2022-12-07T09:19:56.866742Z","shell.execute_reply":"2022-12-07T09:20:01.354877Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# <span><h1 style = \"font-family: garamond; font-size: 40px; font-style: normal; letter-spcaing: 3px; background-color: #f6f5f5; color :#fe346e; border-radius: 100px 100px; text-align:center\">Generator✨<a id=\"Model Architecture\"></a></h1></span>\n\n![unet](https://raw.githubusercontent.com/prabhanjan-jadhav/image-colorization-using-gans/main/images/UNet%20architechture.png)","metadata":{}},{"cell_type":"code","source":"# Utility function required to crop a tensor while passing through skip connection\ndef crop_tensor(tensor, target_tensor):\n    target_size = target_tensor.size()[2]\n    tensor_size = tensor.size()[2]\n    delta = tensor_size - target_size\n    delta = delta // 2\n    return tensor[\n        :,\n        :,\n        delta:tensor_size - delta,\n        delta:tensor_size - delta\n    ]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class UNet(nn.Module):\n    def __init__(self, in_channels=1, out_channels=2):\n        super().__init__()\n        self.dconv1 = nn.Conv2d(in_channels, 64, kernel_size=3) \n        self.dconv2 = nn.Conv2d(64, 128, kernel_size=4)\n        self.dconv3 = nn.Conv2d(128, 256, kernel_size=3)\n        self.dconv4 = nn.Conv2d(256, 512, kernel_size=3)\n        self.dconv5 = nn.Conv2d(512, 512, kernel_size=3)\n\n        self.uconv1 = nn.Conv2d(512, 256, kernel_size=3)\n        self.uconv2 = nn.Conv2d(256, 128, kernel_size=3)\n        self.uconv3 = nn.Conv2d(128, 64, kernel_size=3)\n        self.uconv4 = nn.Conv2d(64, 2, kernel_size=1)\n\n        self.maxpool2d = nn.MaxPool2d(2)\n\n        self.trans1 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)\n        self.trans2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)\n        self.trans3 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)\n\n    def forward(self, image):\n        x1 = self.dconv1(image)\n        x2 = self.maxpool2d(x1) \n\n        x3 = self.dconv2(x2) \n        x4 = self.maxpool2d(x3)\n\n        x5 = self.dconv3(x4)    \n        x6 = self.maxpool2d(x5)\n\n        x7 = self.dconv4(x6)\n        x8 = self.dconv5(x7)\n\n        x9 = self.trans1(x8) \n        y = crop_tensor(x5, x9)\n        x9 = torch.cat([x9, y], axis=1)\n        x10 = self.uconv1(x9)\n\n        x11 = self.trans2(x10)\n\n        y = crop_tensor(x3, x11)\n        x11 = torch.cat([x11, y], axis=1)\n        x12 = self.uconv2(x11)\n\n        x13 = self.trans3(x12)\n        y = crop_tensor(x1, x13)\n        x13 = torch.cat([x13, y], axis=1)\n        x14 = self.uconv3(x13)\n\n        out = self.uconv4(x14)\n#         print(f\"Output image size : {out.size()}\")\n        return out\n    \nif __name__ == \"__main__\":\n    model = UNet()\n    print(summary(model, (1,1,256,256)))        ","metadata":{"id":"RAe9u4APbKDb","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# <span><h1 style = \"font-family: garamond; font-size: 40px; font-style: normal; letter-spcaing: 3px; background-color: #f6f5f5; color :#fe346e; border-radius: 100px 100px; text-align:center\">Discriminator </h1></span>","metadata":{"id":"wSFn50LvbKDf"}},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self):\n        super().__init__()\n\n        self.conv1 = nn.Conv2d(3,64,4,2,1)\n        self.leakyrelu1 = nn.LeakyReLU(0.2, inplace=True)\n\n        self.conv2 = nn.Conv2d(64,128,4,2,1)\n        self.bn1 = nn.BatchNorm2d(128)\n        self.leakyrelu2 = nn.LeakyReLU(0.2, inplace=True)\n\n        self.conv3 = nn.Conv2d(128,256,4,2,1)\n        self.bn2 = nn.BatchNorm2d(256)\n        self.leakyrelu3 = nn.LeakyReLU(0.2, inplace=True)\n\n        self.conv4 = nn.Conv2d(256,512,4,1,1)\n        self.bn3 = nn.BatchNorm2d(512)\n        self.leakyrelu4 = nn.LeakyReLU(0.2, inplace=True)\n\n        self.conv5 = nn.Conv2d(512,1,4,1,1)\n\n    def forward(self, x):\n        layers = [self.conv1, self.leakyrelu1, self.conv2, self.bn1, self.leakyrelu2, self.conv3, \n                  self.bn2, self.leakyrelu3, self.conv4, self.bn3, self.leakyrelu4, self.conv5]\n\n        model = nn.Sequential(*layers)\n        self.model = nn.Sequential(*model)\n        return model(x)","metadata":{"id":"kHM1-6VabKDf","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Let's take a look at its blocks:","metadata":{"id":"y24XzZK5bKDg"}},{"cell_type":"code","source":"summary(Discriminator())","metadata":{"id":"0reri1LybKDg","outputId":"a5cd5c16-ec0f-40c1-ec97-0ba2357c47fa","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"And its output shape:","metadata":{"id":"6pLthOgybKDg"}},{"cell_type":"code","source":"discriminator = Discriminator()\ndummy_input = torch.randn(16, 3, 194, 194) # batch_size, channels, size, size\nout = discriminator(dummy_input)\nprint(out.shape)","metadata":{"id":"9RGmmBHgbKDg","outputId":"4519e281-731b-4ff0-e08b-61b53220fa39","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# <span><h1 style = \"font-family: garamond; font-size: 40px; font-style: normal; letter-spcaing: 3px; background-color: #f6f5f5; color :#fe346e; border-radius: 100px 100px; text-align:center\">Loss Function 🧭<a id=\"Loss function\"></a></h1></span>\n* **Loss function we optimize** :\n\\begin{equation}\nG^* = arg min_G max_D L_{cGAN} (G,D) + \\lambda L_{L1} (G)\n\\end{equation}\n\n\n* **L1 Loss** : \n\\begin{equation}\nL_{L1} (G) = \\mathbb{E}_{x,y,z}[\\|y-G(x,z)\\|_1]\n\\end{equation}\n\n\n* **GAN Loss** : \n\\begin{equation}\nL_{cGAN}(G,D) = \\mathbb{E}_{x,y}[logD(x,y)] + \\mathbb{E}_{x,z}[log(1-D(x, G(x,z))]\n\\end{equation}\nWhere **x** as the grayscale image, **z** as the input noise for the generator, and **y** as the 2-channel output we want from the generator (it can also represent the 2 color channels of a real image). Also, **G** is the generator model and **D** is the discriminator. \n","metadata":{}},{"cell_type":"code","source":"class GANLoss(nn.Module):\n    def __init__(self, gan_mode='vanilla', real_label=1.0, fake_label=0.0):\n        super().__init__()\n        # If you have parameters in your model, which should be saved and restored\n        # in the state_dict, but not trained by the optimizer, you should register them as buffers.\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":{"id":"W-Dr8X_sbKDh","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# <span><h1 style = \"font-family: garamond; font-size: 40px; font-style: normal; letter-spcaing: 3px; background-color: #f6f5f5; color :#fe346e; border-radius: 100px 100px; text-align:center\">Model Initialization⚙️</h1></span>","metadata":{}},{"cell_type":"code","source":"def 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":{"id":"oZRvVRoMbKDi","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# <span><h1 style = \"font-family: garamond; font-size: 40px; font-style: normal; letter-spcaing: 3px; background-color: #f6f5f5; color :#fe346e; border-radius: 100px 100px; text-align:center\">Model🪄</h1></span>","metadata":{"id":"QkuV0DiibKDi"}},{"cell_type":"code","source":"class MainModel(nn.Module):\n    def __init__(self, net_G=None, lr_G=2e-4, lr_D=2e-4, \n                 beta1=0.5, beta2=0.999, lambda_L1=100.):\n        super().__init__()\n        \n        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n        self.lambda_L1 = lambda_L1\n        \n        if net_G is None:\n            self.net_G = init_model(UNet(), self.device)\n        else:\n            self.net_G = net_G.to(self.device)\n        self.net_D = init_model(Discriminator(), self.device)\n        self.GANcriterion = GANLoss(gan_mode='vanilla').to(self.device)\n        self.L1criterion = nn.L1Loss()\n        self.opt_G = optim.Adam(self.net_G.parameters(), lr=lr_G, betas=(beta1, beta2))\n        self.opt_D = optim.Adam(self.net_D.parameters(), lr=lr_D, betas=(beta1, beta2))\n    \n    def set_requires_grad(self, model, requires_grad=True):\n        for p in model.parameters():\n            p.requires_grad = requires_grad\n        \n    def setup_input(self, data):\n        self.L = data['L'].to(self.device)\n        self.ab = data['ab'].to(self.device)\n        img = torch.cat((self.L, self.ab), 1)\n        img = transforms.Resize((194, 194), Image.BICUBIC)(img)\n        self.L_resized = img[:, [0], ...]\n        \n    def forward(self):\n        self.fake_color = self.net_G(self.L)\n        img = torch.cat((self.L, self.fake_color),1)\n        img = transforms.Resize((256, 256), Image.BICUBIC)(img)\n        self.fake_color = img[:, [1,2], ...]\n    \n    def backward_D(self, data):\n        fake_image = torch.cat([self.L, self.fake_color], dim=1) # concat side by side\n        fake_preds = self.net_D(fake_image.detach()) # detach() : It returns a new tensor that doesn't require a gradient.\n        self.loss_D_fake = self.GANcriterion(fake_preds, False)\n        \n        real_image = torch.cat([self.L, self.ab], dim=1)\n        real_preds = self.net_D(real_image)\n        self.loss_D_real = self.GANcriterion(real_preds, True)\n        self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5\n        self.loss_D.backward()\n    \n    def backward_G(self, data):\n        fake_image = torch.cat([self.L, self.fake_color], dim=1)\n        fake_preds = self.net_D(fake_image)\n        self.loss_G_GAN = self.GANcriterion(fake_preds, True)\n\n        self.loss_G_L1 = self.L1criterion(self.fake_color, 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.net_D.train()\n        self.set_requires_grad(self.net_D, True)\n        self.opt_D.zero_grad()\n        self.backward_D(data)\n        self.opt_D.step()\n        \n        self.net_G.train()\n        self.set_requires_grad(self.net_D, False)\n        self.opt_G.zero_grad()\n        self.backward_G(data)\n        self.opt_G.step()","metadata":{"id":"DqaCPsCGbKDi","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# <span><h1 style = \"font-family: garamond; font-size: 40px; font-style: normal; letter-spcaing: 3px; background-color: #f6f5f5; color :#fe346e; border-radius: 100px 100px; text-align:center\">Utility functions🎈</h1></span>","metadata":{"id":"gd6bvRvhbKDj"}},{"cell_type":"code","source":"class AverageMeter:\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    loss_D_fake = AverageMeter()\n    loss_D_real = AverageMeter()\n    loss_D = AverageMeter()\n    loss_G_GAN = AverageMeter()\n    loss_G_L1 = AverageMeter()\n    loss_G = AverageMeter()\n    \n    return {'loss_D_fake': loss_D_fake,\n            'loss_D_real': loss_D_real,\n            'loss_D': loss_D,\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\ndef lab_to_rgb(L, ab):\n    \"\"\"\n    Takes a batch of images\n    \"\"\"\n    \n    L = (L + 1.) * 50.\n    ab = ab * 110.\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=False):\n    model.net_G.eval()\n    with torch.no_grad():\n        model.setup_input(data)\n        model.forward()\n    model.net_G.train()\n    fake_color = model.fake_color.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}\")","metadata":{"id":"w8miTDYZbKDj","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# <span><h1 style = \"font-family: garamond; font-size: 40px; font-style: normal; letter-spcaing: 3px; background-color: #f6f5f5; color :#fe346e; border-radius: 100px 100px; text-align:center\">Training💪</h1></span>","metadata":{}},{"cell_type":"markdown","source":"Now to train the model run the below cell. Running for about 20 epochs, taking approx. 5 mins per epoch, would give reasonable results. \\\nInstead, you can directly download the weights for the model trained for 20 epochs in the following cell.","metadata":{}},{"cell_type":"code","source":"def train_model(model, train_dl, epochs, display_every=200):\n    data = next(iter(val_dl)) # getting a batch for visualizing the model output after fixed intrvals\n    for e in range(epochs):\n        loss_meter_dict = create_loss_meters() # function returing a dictionary of objects to \n        i = 0                                  # log the losses of the complete network\n        for data in tqdm(train_dl):\n            model.setup_input(data) \n            model.optimize()\n            update_losses(model, loss_meter_dict, count=data['L'].size(0)) # function updating the log objects\n            i += 1\n            if i % display_every == 0:\n                print(f\"\\nEpoch {e+1}/{epochs}\")\n                print(f\"Iteration {i}/{len(train_dl)}\")\n                log_results(loss_meter_dict) # function to print out the losses\n                visualize(model, data, save=False) # function displaying the model's outputs\n        if e%5==0:\n            torch.save(model.state_dict(), f\"/kaggle/working/model_{e+1}.pt\")\n            \nmodel = MainModel()\n# train_model(model, train_dl, 20)","metadata":{"id":"NPmcqCczbKDk","scrolled":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# downloading weights for the model\n!gdown --id 1owvzniVc_PQ3xqNHscGlR_mrbpdy7e_G","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# create the model object and load the weights\nmodel = MainModel()\nmodel.load_state_dict(torch.load('/kaggle/working/model_21.pt'))","metadata":{"id":"NPmcqCczbKDk","scrolled":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Visualize the model\nmodel.eval()","metadata":{"id":"NPmcqCczbKDk","scrolled":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Generating output on validation dataset\nfor data in tqdm(val_dl):\n    model.setup_input(data)\n    model.optimize()\n    visualize(model, data, save=False)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"This is the end of the 1st way of implementation. \nThe following section contains a modified implementation which will improve the results.","metadata":{}},{"cell_type":"markdown","source":"<h1 style = \"font-size:60px; font-family:Garamond ; font-weight : normal; background-color: #f6f5f5 ; color : #fe346e; text-align: center; border-radius: 100px 100px;\">A Sophisticated Approach🚀<a id=\"Another approach\"></a></h1>\nThis module uses fastai and torchvision models to build the generator model. \n\nA ResNet pretrained on ImageNet classification task is added as a backbone to the UNet of the generator. The discriminator is the same as before.","metadata":{"id":"pSgjm619bKDl"}},{"cell_type":"code","source":"# !pip install fastai==2.4\nfrom fastai.vision.learner import create_body\nfrom torchvision.models.resnet import resnet18\nfrom fastai.vision.models.unet import DynamicUnet","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def build_res_unet(n_input=1, n_output=2, size=256):\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    body = create_body(resnet18(), pretrained=True, n_in=n_input, cut=-2)\n    net_G = DynamicUnet(body, n_output, (size, size)).to(device)\n    return net_G","metadata":{"id":"mBa0FHRlbKDm","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def pretrain_generator(net_G, train_dl, opt, criterion, epochs):\n    ''' Pretraining generator on image colorization task using L1 loss.\n    ResNet backbone has pretrained weights'''\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 = net_G(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}\")\n\nnet_G = build_res_unet(n_input=1, n_output=2, size=256)\nopt = optim.Adam(net_G.parameters(), lr=1e-4)\ncriterion = nn.L1Loss()  \npretrain_generator(net_G, train_dl, opt, criterion, 20)\ntorch.save(net_G.state_dict(), \"res18-unet.pt\")","metadata":{"id":"RREGAO_MbKDm","outputId":"3ad39156-016e-42e5-be47-ebd3eb58f523","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"'''\nTraining the entire model for 20 epochs\n'''\nnet_G = build_res_unet(n_input=1, n_output=2, size=256)\nnet_G.load_state_dict(torch.load(\"res18-unet.pt\", map_location=device))\nmodel = MainModel(net_G=net_G)\ntrain_model(model, train_dl, 20)\ntorch.save(model.state_dict(), \"final_model_weights.pt\")","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"''' Download the weights for trained model\n'''\n!gdown --id 1UY5a07bVofwAyV7rI8konccnFCJjIKPZ\n!gdown --id 1lR6DcS4m5InSbZ5y59zkH2mHt_4RQ2KV\n\n'''\n\n    Build the model object and initialize the weights\n\n'''\nnet_G = build_res_unet(n_input=1, n_output=2, size=256)\nnet_G.load_state_dict(torch.load(\"/kaggle/working/res18-unet.pt\", map_location=device))\nmodel = MainModel(net_G=net_G)\nmodel.load_state_dict(torch.load(\"/kaggle/working/final_model_weights.pt\", map_location=device))","metadata":{"id":"JSOueoq2bKDo","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"'''\nVisualize the outputs on validation dataset\n\n'''\nfor data in tqdm(val_dl):\n    model.setup_input(data)\n    model.optimize()\n    visualize(model, data, save=False)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"'''\nGenerate output on your own set of images\n\n'''\nimg_path = '/path/to/folder'\nprint(img_path)\npaths = glob.glob(img_path + \"/*\")\nidxs = np.arange(len(paths))\n\ntest_dl = make_dataloaders(paths=paths, split='val')\nfor data in tqdm(test_dl):\n  model.setup_input(data)\n  model.optimize()\n  visualize(model, data, save=False)","metadata":{},"execution_count":null,"outputs":[]}]}