image-colorization-starter.ipynb

{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Image Colourisation using Autoencoders"},{"metadata":{},"cell_type":"markdown","source":"## Imports\nImporting all the necessary libraries required to run the following code for image colourization."},{"metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","execution":{"iopub.execute_input":"2020-08-30T16:53:38.593966Z","iopub.status.busy":"2020-08-30T16:53:38.593021Z","iopub.status.idle":"2020-08-30T16:53:38.594926Z","shell.execute_reply":"2020-08-30T16:53:38.595443Z"},"papermill":{"duration":0.022598,"end_time":"2020-08-30T16:53:38.595602","exception":false,"start_time":"2020-08-30T16:53:38.573004","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"import numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport matplotlib.pyplot as plt # visualising and plotting results\n\nimport os","execution_count":null,"outputs":[]},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:53:38.649161Z","iopub.status.busy":"2020-08-30T16:53:38.648444Z","iopub.status.idle":"2020-08-30T16:53:40.310655Z","shell.execute_reply":"2020-08-30T16:53:40.310102Z"},"papermill":{"duration":1.679608,"end_time":"2020-08-30T16:53:40.310780","exception":false,"start_time":"2020-08-30T16:53:38.631172","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"from torchvision.datasets import ImageFolder\nfrom torch.utils.data import DataLoader\nfrom torch.utils.data.dataset import random_split\n\nimport torchvision.transforms as T\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom tqdm.notebook import tqdm\nfrom skimage.color import rgb2lab, lab2rgb, rgb2gray","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.011725,"end_time":"2020-08-30T16:53:38.619321","exception":false,"start_time":"2020-08-30T16:53:38.607596","status":"completed"},"tags":[]},"cell_type":"markdown","source":"## Loading the Data"},{"metadata":{"papermill":{"duration":0.012125,"end_time":"2020-08-30T16:53:40.337059","exception":false,"start_time":"2020-08-30T16:53:40.324934","status":"completed"},"tags":[]},"cell_type":"markdown","source":"Defining the directory for the dataset to be used."},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","execution":{"iopub.execute_input":"2020-08-30T16:53:40.365779Z","iopub.status.busy":"2020-08-30T16:53:40.365090Z","iopub.status.idle":"2020-08-30T16:53:40.368922Z","shell.execute_reply":"2020-08-30T16:53:40.369363Z"},"papermill":{"duration":0.020353,"end_time":"2020-08-30T16:53:40.369497","exception":false,"start_time":"2020-08-30T16:53:40.349144","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"DATA_DIR = '../input/imagenet'","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.011642,"end_time":"2020-08-30T16:53:40.393706","exception":false,"start_time":"2020-08-30T16:53:40.382064","status":"completed"},"tags":[]},"cell_type":"markdown","source":"Loading images contained in the data directory into a avriable called dataset using ImageFolder, and applying two transforms to all the images:\n1. Resizing the non-uniformly sized images into 256x256 images.\n2. Converting them into tensors."},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:53:40.422312Z","iopub.status.busy":"2020-08-30T16:53:40.421733Z","iopub.status.idle":"2020-08-30T16:55:57.374381Z","shell.execute_reply":"2020-08-30T16:55:57.374907Z"},"papermill":{"duration":136.969334,"end_time":"2020-08-30T16:55:57.375087","exception":false,"start_time":"2020-08-30T16:53:40.405753","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"dataset = ImageFolder(DATA_DIR, transform=T.Compose([T.Resize((256, 256)),T.ToTensor()]))","execution_count":null,"outputs":[]},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:55:57.407495Z","iopub.status.busy":"2020-08-30T16:55:57.406147Z","iopub.status.idle":"2020-08-30T16:55:57.410157Z","shell.execute_reply":"2020-08-30T16:55:57.410637Z"},"papermill":{"duration":0.022841,"end_time":"2020-08-30T16:55:57.410753","exception":false,"start_time":"2020-08-30T16:55:57.387912","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"len(dataset)","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.011767,"end_time":"2020-08-30T16:55:57.434598","exception":false,"start_time":"2020-08-30T16:55:57.422831","status":"completed"},"tags":[]},"cell_type":"markdown","source":"Creating the training and validation sets. A random seed value is set to have the same training and validation datasets each time the notebook is run, and the two datasets are split using the random_split function from the pytorch libraries."},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:55:57.464116Z","iopub.status.busy":"2020-08-30T16:55:57.463331Z","iopub.status.idle":"2020-08-30T16:55:57.488110Z","shell.execute_reply":"2020-08-30T16:55:57.487578Z"},"papermill":{"duration":0.041667,"end_time":"2020-08-30T16:55:57.488200","exception":false,"start_time":"2020-08-30T16:55:57.446533","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"random_seed = 42\ntorch.manual_seed(random_seed)\n\nval_size = 1000\ntrain_size = len(dataset) - val_size\n\ntrain_ds, val_ds = random_split(dataset, [train_size, val_size])\nlen(train_ds), len(val_ds)","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.011794,"end_time":"2020-08-30T16:55:57.512171","exception":false,"start_time":"2020-08-30T16:55:57.500377","status":"completed"},"tags":[]},"cell_type":"markdown","source":"The batch size is set to be 128."},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:55:57.539764Z","iopub.status.busy":"2020-08-30T16:55:57.539150Z","iopub.status.idle":"2020-08-30T16:55:57.542913Z","shell.execute_reply":"2020-08-30T16:55:57.543344Z"},"papermill":{"duration":0.019311,"end_time":"2020-08-30T16:55:57.543456","exception":false,"start_time":"2020-08-30T16:55:57.524145","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"batch_size = 128","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.011753,"end_time":"2020-08-30T16:55:57.567201","exception":false,"start_time":"2020-08-30T16:55:57.555448","status":"completed"},"tags":[]},"cell_type":"markdown","source":"Loading the training and validation datasets into the CPU using DataLoader."},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:55:57.596955Z","iopub.status.busy":"2020-08-30T16:55:57.596080Z","iopub.status.idle":"2020-08-30T16:55:57.598345Z","shell.execute_reply":"2020-08-30T16:55:57.598791Z"},"papermill":{"duration":0.019423,"end_time":"2020-08-30T16:55:57.598904","exception":false,"start_time":"2020-08-30T16:55:57.579481","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"train_loader = DataLoader(train_ds, batch_size = batch_size, shuffle = True, num_workers = 4, pin_memory = True)\nval_loader = DataLoader(val_ds, batch_size = batch_size, num_workers = 4, pin_memory = True)","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.011768,"end_time":"2020-08-30T16:55:57.622750","exception":false,"start_time":"2020-08-30T16:55:57.610982","status":"completed"},"tags":[]},"cell_type":"markdown","source":"Use CUDA if available else fall back to CPU for training."},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:55:57.984898Z","iopub.status.busy":"2020-08-30T16:55:57.984082Z","iopub.status.idle":"2020-08-30T16:55:57.987364Z","shell.execute_reply":"2020-08-30T16:55:57.987866Z"},"papermill":{"duration":0.353214,"end_time":"2020-08-30T16:55:57.988000","exception":false,"start_time":"2020-08-30T16:55:57.634786","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\ndevice","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.012162,"end_time":"2020-08-30T16:55:58.012974","exception":false,"start_time":"2020-08-30T16:55:58.000812","status":"completed"},"tags":[]},"cell_type":"markdown","source":"Define a function to load data into the device assigned above."},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:55:58.043295Z","iopub.status.busy":"2020-08-30T16:55:58.042421Z","iopub.status.idle":"2020-08-30T16:55:58.045478Z","shell.execute_reply":"2020-08-30T16:55:58.044997Z"},"papermill":{"duration":0.020358,"end_time":"2020-08-30T16:55:58.045600","exception":false,"start_time":"2020-08-30T16:55:58.025242","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"def to_device(data, device):\n    if isinstance(data, (list, tuple)):\n        return [to_device(x, device) for x in data]\n    return data.to(device, non_blocking = True)","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.012074,"end_time":"2020-08-30T16:55:58.069990","exception":false,"start_time":"2020-08-30T16:55:58.057916","status":"completed"},"tags":[]},"cell_type":"markdown","source":"Define a custom DataLoader to load data onto the available device."},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:55:58.101770Z","iopub.status.busy":"2020-08-30T16:55:58.100933Z","iopub.status.idle":"2020-08-30T16:55:58.103191Z","shell.execute_reply":"2020-08-30T16:55:58.103768Z"},"papermill":{"duration":0.02122,"end_time":"2020-08-30T16:55:58.103884","exception":false,"start_time":"2020-08-30T16:55:58.082664","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"class DeviceDataLoader():\n    def __init__(self, dl, device):\n        self.dl = dl\n        self.device = device\n\n    def __iter__(self):\n        for batch in self.dl:\n            yield to_device(batch, self.device)\n  \n    def __len__(self):\n        return len(self.dl)","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.011824,"end_time":"2020-08-30T16:55:58.127969","exception":false,"start_time":"2020-08-30T16:55:58.116145","status":"completed"},"tags":[]},"cell_type":"markdown","source":"Wrapping the training and validation dataloaders in our custom and efficient `DeviceDataLoader`."},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:55:58.157789Z","iopub.status.busy":"2020-08-30T16:55:58.157011Z","iopub.status.idle":"2020-08-30T16:55:58.159466Z","shell.execute_reply":"2020-08-30T16:55:58.160058Z"},"papermill":{"duration":0.020052,"end_time":"2020-08-30T16:55:58.160184","exception":false,"start_time":"2020-08-30T16:55:58.140132","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"train_loader = DeviceDataLoader(train_loader, device)\nval_loader = DeviceDataLoader(val_loader, device)","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.011903,"end_time":"2020-08-30T16:55:58.186027","exception":false,"start_time":"2020-08-30T16:55:58.174124","status":"completed"},"tags":[]},"cell_type":"markdown","source":"## Defining the Models\nThe images are originally using the RGB colormodel. But for the purpose of image colourization, they have to be converted to Lab. The CIELAB is a colorspace which has 3 channels: L, a, and b. The function defined below converts images from RGB to CIELAB, and then splits the L and ab channels into the variables X and Y respectively. The range of ab channels is from -128 to 127, hence, the ab channels are normalized by dividing it by 128. The function returns finally returns X and Y loaded into cuda."},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:55:58.220795Z","iopub.status.busy":"2020-08-30T16:55:58.220023Z","iopub.status.idle":"2020-08-30T16:55:58.223238Z","shell.execute_reply":"2020-08-30T16:55:58.222589Z"},"papermill":{"duration":0.023905,"end_time":"2020-08-30T16:55:58.223356","exception":false,"start_time":"2020-08-30T16:55:58.199451","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"def generate_l_ab(images): \n    lab = rgb2lab(images.permute(0, 2, 3, 1).cpu().numpy())\n    X = lab[:,:,:,0]\n    X = X.reshape(X.shape+(1,))\n    Y = lab[:,:,:,1:] / 128\n    return to_device(torch.tensor(X, dtype = torch.float).permute(0, 3, 1, 2), device),to_device(torch.tensor(Y, dtype = torch.float).permute(0, 3, 1, 2), device)","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.015253,"end_time":"2020-08-30T16:55:58.252155","exception":false,"start_time":"2020-08-30T16:55:58.236902","status":"completed"},"tags":[]},"cell_type":"markdown","source":"The Base Model class is defined as containing three functions. \n1. The *training_batch* function takes the batch of 128 images as input, generates the values of L and ab channels using the *generate_l_ab* function, gets the predicted ab channels using the forward function of the model, and calculates the MSE Loss between the actual and predicted ab channels.\n2. The *validation_batch* functions performs the same task as the *training_batch* function, except for the images in the validation dataset, which is evident from the function names.\n3. The function *validation_end_epoch* returns the average loss on the validation dataset."},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:55:58.289943Z","iopub.status.busy":"2020-08-30T16:55:58.289136Z","iopub.status.idle":"2020-08-30T16:55:58.292133Z","shell.execute_reply":"2020-08-30T16:55:58.291644Z"},"papermill":{"duration":0.026257,"end_time":"2020-08-30T16:55:58.292230","exception":false,"start_time":"2020-08-30T16:55:58.265973","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"class BaseModel(nn.Module):\n    def training_batch(self, batch):\n        images, _ = batch\n        X, Y = generate_l_ab(images)\n        outputs = self.forward(X)\n        loss = F.mse_loss(outputs, Y)\n        return loss\n\n    def validation_batch(self, batch):\n        images, _ = batch\n        X, Y = generate_l_ab(images)\n        outputs = self.forward(X)\n        loss = F.mse_loss(outputs, Y)\n        return {'val_loss' : loss.item()}\n\n    def validation_end_epoch(self, outputs):\n        epoch_loss = sum([x['val_loss'] for x in outputs]) / len(outputs)\n        return {'epoch_loss' : epoch_loss}","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.011838,"end_time":"2020-08-30T16:55:58.316828","exception":false,"start_time":"2020-08-30T16:55:58.304990","status":"completed"},"tags":[]},"cell_type":"markdown","source":"A helper function is defined to get the appropriate padding in order to keep the size of the output image same as the input image during convolution."},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:55:58.347486Z","iopub.status.busy":"2020-08-30T16:55:58.346746Z","iopub.status.idle":"2020-08-30T16:55:58.349698Z","shell.execute_reply":"2020-08-30T16:55:58.349218Z"},"papermill":{"duration":0.020118,"end_time":"2020-08-30T16:55:58.349793","exception":false,"start_time":"2020-08-30T16:55:58.329675","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"def get_padding(kernel_size: int, stride: int = 1, dilation: int = 1, **_) -> int:\n    padding = ((stride - 1) + dilation * (kernel_size - 1)) // 2\n    return padding","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.011887,"end_time":"2020-08-30T16:55:58.374468","exception":false,"start_time":"2020-08-30T16:55:58.362581","status":"completed"},"tags":[]},"cell_type":"markdown","source":"The `Encoder_Decoder_v1` class is an extension of the Base Model and it contains the code for the version 1 architecture of the model."},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:55:58.419324Z","iopub.status.busy":"2020-08-30T16:55:58.408494Z","iopub.status.idle":"2020-08-30T16:55:58.421471Z","shell.execute_reply":"2020-08-30T16:55:58.422062Z"},"papermill":{"duration":0.034922,"end_time":"2020-08-30T16:55:58.422175","exception":false,"start_time":"2020-08-30T16:55:58.387253","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"class Encoder_Decoder_v1(BaseModel):\n    def __init__(self):\n        super().__init__()\n        self.network = nn.Sequential(\n            nn.Conv2d(1, 64, kernel_size = 3, stride = 2, padding = get_padding(3, 2)),\n            nn.ReLU(),\n            nn.BatchNorm2d(64),\n            nn.Conv2d(64, 128, kernel_size = 3, padding = get_padding(3)),\n            nn.ReLU(),\n            nn.BatchNorm2d(128),\n        \n            nn.Conv2d(128, 128, kernel_size = 3, stride = 2, padding = get_padding(3, 2)),\n            nn.ReLU(),\n            nn.BatchNorm2d(128),\n            nn.Conv2d(128, 256, kernel_size = 3, padding = get_padding(3)),\n            nn.ReLU(),\n            nn.BatchNorm2d(256),\n        \n            nn.Conv2d(256, 256, kernel_size = 3, stride = 2, padding = get_padding(3, 2)),\n            nn.ReLU(),\n            nn.BatchNorm2d(256),\n            nn.Conv2d(256, 512, kernel_size = 3, padding = get_padding(3)),\n            nn.ReLU(),\n            nn.BatchNorm2d(512),\n            \n            nn.Conv2d(512, 512, kernel_size = 3, padding = get_padding(3)),\n            nn.ReLU(),\n            nn.BatchNorm2d(512),\n            nn.Conv2d(512, 256, kernel_size = 3, padding = get_padding(3)),\n            nn.ReLU(),\n            nn.BatchNorm2d(256),\n        \n            nn.Conv2d(256, 128, kernel_size = 3, padding = get_padding(3)),\n            nn.Upsample(size = (64,64)),\n            nn.Conv2d(128, 64, kernel_size = 3, padding = get_padding(3)),\n            nn.Upsample(size = (128,128)),\n            nn.Conv2d(64, 32, kernel_size = 3, padding = get_padding(3)),\n            nn.Conv2d(32, 16, kernel_size = 3, padding = get_padding(3)),\n            nn.Conv2d(16, 2, kernel_size = 3, padding = get_padding(3)),\n            nn.Tanh(),\n            nn.Upsample(size = (256,256))\n    )\n\n    def forward(self, images):\n        return self.network(images)     \n    ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The `Encoder_Decoder_v2` class is an extension of the Base Model and it contains the code for the version 2 architecture of the model which uses ResNet18 for the Encoder part."},{"metadata":{"trusted":true},"cell_type":"code","source":"class Encoder_Decoder_v2(Base_Model):\n    def __init__(self, input_size=128):\n        super(Encoder_Decoder, self).__init__()\n        MIDLEVEL_FEATURE_SIZE = 128\n        \n        \n        # Encoder Part\n        resnet = models.resnet18(num_classes=365) \n        resnet.conv1.weight = nn.Parameter(resnet.conv1.weight.sum(dim=1).unsqueeze(1)) \n        self.midlevel_resnet = nn.Sequential(*list(resnet.children())[0:6]) # Extracting the first\n                                                                            # 6 convolutional layers\n        \n        # Decoder Part\n        self.upsample = nn.Sequential(     \n          nn.Conv2d(MIDLEVEL_FEATURE_SIZE, 128, kernel_size=3, stride=1, padding=1),\n          nn.BatchNorm2d(128),\n          nn.ReLU(),\n          nn.Upsample(scale_factor=2),\n          nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1),\n          nn.BatchNorm2d(64),\n          nn.ReLU(),\n          nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),\n          nn.BatchNorm2d(64),\n          nn.ReLU(),\n          nn.Upsample(scale_factor=2),\n          nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1),\n          nn.BatchNorm2d(32),\n          nn.ReLU(),\n          nn.Conv2d(32, 2, kernel_size=3, stride=1, padding=1),\n          nn.Upsample(scale_factor=2)\n        )\n        \n    def forward(self, input):\n\n        # Pass input through ResNet-gray to extract features\n        midlevel_features = self.midlevel_resnet(input)\n\n        # Upsample to get colors\n        output = self.upsample(midlevel_features)\n        return output","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.011848,"end_time":"2020-08-30T16:55:58.446094","exception":false,"start_time":"2020-08-30T16:55:58.434246","status":"completed"},"tags":[]},"cell_type":"markdown","source":"The model is defined and loaded to cuda."},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:55:58.483145Z","iopub.status.busy":"2020-08-30T16:55:58.482467Z","iopub.status.idle":"2020-08-30T16:56:05.113706Z","shell.execute_reply":"2020-08-30T16:56:05.113180Z"},"papermill":{"duration":6.655535,"end_time":"2020-08-30T16:56:05.113821","exception":false,"start_time":"2020-08-30T16:55:58.458286","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"model = Encoder_Decoder_v1()\nto_device(model, device)","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.012117,"end_time":"2020-08-30T16:56:05.138712","exception":false,"start_time":"2020-08-30T16:56:05.126595","status":"completed"},"tags":[]},"cell_type":"markdown","source":"## Training the Model\nThe *validate* and *fit* functions are defined to keep track of the loss and train the model."},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:56:05.174768Z","iopub.status.busy":"2020-08-30T16:56:05.173993Z","iopub.status.idle":"2020-08-30T16:56:05.176460Z","shell.execute_reply":"2020-08-30T16:56:05.176998Z"},"papermill":{"duration":0.026016,"end_time":"2020-08-30T16:56:05.177107","exception":false,"start_time":"2020-08-30T16:56:05.151091","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"@torch.no_grad()\ndef validate(model, val_loader):\n    model.eval()\n    outputs = [model.validation_batch(batch) for batch in val_loader]\n    return model.validation_end_epoch(outputs)\n\ndef fit(model, epochs, learning_rate, train_loader, val_loader, optimization_func = torch.optim.SGD):\n    torch.cuda.empty_cache()\n    history = []\n    optimizer = optimization_func(model.parameters(), learning_rate)\n    for epoch in range(epochs):\n        train_losses = []\n        model.train()\n        for batch in tqdm(train_loader):\n            loss = model.training_batch(batch)\n            train_losses.append(loss)\n            loss.backward()\n            optimizer.step()\n            optimizer.zero_grad()\n        result = validate(model, val_loader)\n        result['train_loss'] = torch.stack(train_losses).mean().item()\n        history.append(result)\n        print('Epoch: {}, Train loss: {:.4f}, Validation loss: {:.4f}'.format(epoch, result['train_loss'], result['epoch_loss']))\n    return history","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.012172,"end_time":"2020-08-30T16:56:05.201578","exception":false,"start_time":"2020-08-30T16:56:05.189406","status":"completed"},"tags":[]},"cell_type":"markdown","source":"Initializing the loss."},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:56:05.229643Z","iopub.status.busy":"2020-08-30T16:56:05.228878Z","iopub.status.idle":"2020-08-30T16:56:05.231719Z","shell.execute_reply":"2020-08-30T16:56:05.231229Z"},"papermill":{"duration":0.017949,"end_time":"2020-08-30T16:56:05.231815","exception":false,"start_time":"2020-08-30T16:56:05.213866","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"#history = [validate(model, val_loader)]\n#history","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.012338,"end_time":"2020-08-30T16:56:05.258730","exception":false,"start_time":"2020-08-30T16:56:05.246392","status":"completed"},"tags":[]},"cell_type":"markdown","source":"Training."},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:56:05.286901Z","iopub.status.busy":"2020-08-30T16:56:05.286152Z","iopub.status.idle":"2020-08-30T16:56:05.289117Z","shell.execute_reply":"2020-08-30T16:56:05.288658Z"},"papermill":{"duration":0.01816,"end_time":"2020-08-30T16:56:05.289214","exception":false,"start_time":"2020-08-30T16:56:05.271054","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"#history += fit(model, 5, 0.001, train_loader, val_loader, torch.optim.Adam)","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.011959,"end_time":"2020-08-30T16:56:05.313449","exception":false,"start_time":"2020-08-30T16:56:05.301490","status":"completed"},"tags":[]},"cell_type":"markdown","source":"Saving the model to be able to continue training the model from the same point."},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:56:05.341936Z","iopub.status.busy":"2020-08-30T16:56:05.341297Z","iopub.status.idle":"2020-08-30T16:56:05.378135Z","shell.execute_reply":"2020-08-30T16:56:05.377647Z"},"papermill":{"duration":0.05238,"end_time":"2020-08-30T16:56:05.378233","exception":false,"start_time":"2020-08-30T16:56:05.325853","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"torch.save(model.state_dict(), 'test.pth')","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.012031,"end_time":"2020-08-30T16:56:05.402735","exception":false,"start_time":"2020-08-30T16:56:05.390704","status":"completed"},"tags":[]},"cell_type":"markdown","source":"Loading the pth file to continue training from the same parameters onwards."},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:56:05.431868Z","iopub.status.busy":"2020-08-30T16:56:05.431116Z","iopub.status.idle":"2020-08-30T16:56:05.433600Z","shell.execute_reply":"2020-08-30T16:56:05.434118Z"},"papermill":{"duration":0.019203,"end_time":"2020-08-30T16:56:05.434228","exception":false,"start_time":"2020-08-30T16:56:05.415025","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"def load_checkpoint(filepath): \n    model = Encoder_Decoder_v1()\n    model.load_state_dict(torch.load(filepath))\n    \n    return model","execution_count":null,"outputs":[]},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:56:05.463335Z","iopub.status.busy":"2020-08-30T16:56:05.462561Z","iopub.status.idle":"2020-08-30T16:56:05.935182Z","shell.execute_reply":"2020-08-30T16:56:05.935687Z"},"papermill":{"duration":0.489384,"end_time":"2020-08-30T16:56:05.935823","exception":false,"start_time":"2020-08-30T16:56:05.446439","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"model = load_checkpoint('../input/landscape-test/Landscapes.pth')\nto_device(model, device)","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.012453,"end_time":"2020-08-30T16:56:05.961256","exception":false,"start_time":"2020-08-30T16:56:05.948803","status":"completed"},"tags":[]},"cell_type":"markdown","source":"## Testing the Model\nCombining the L channel and the predicted ab channels and converting it to RGB to obtain the final colour image."},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:56:05.994314Z","iopub.status.busy":"2020-08-30T16:56:05.993451Z","iopub.status.idle":"2020-08-30T16:56:05.995890Z","shell.execute_reply":"2020-08-30T16:56:05.996332Z"},"papermill":{"duration":0.022663,"end_time":"2020-08-30T16:56:05.996449","exception":false,"start_time":"2020-08-30T16:56:05.973786","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"def to_rgb(grayscale_input, ab_output):\n    color_image = torch.cat((grayscale_input, ab_output), 0).numpy() # combine channels\n    print(color_image.shape)\n    color_image = color_image.transpose((1, 2, 0))  # rescale for matplotlib\n    color_image[:, :, 0:1] = color_image[:, :, 0:1]\n    color_image[:, :, 1:3] = (color_image[:, :, 1:3]) * 128\n    color_image = lab2rgb(color_image.astype(np.float64))\n    grayscale_input = grayscale_input.squeeze().numpy()\n    return color_image","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.012475,"end_time":"2020-08-30T16:56:06.021725","exception":false,"start_time":"2020-08-30T16:56:06.009250","status":"completed"},"tags":[]},"cell_type":"markdown","source":"Given a black and white image with 3 channels(R==G==B), converting it to Lab and giving the L channel as input to predict the ab channels and finally, obtaining the predicted coloured version of the image."},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:56:06.054012Z","iopub.status.busy":"2020-08-30T16:56:06.053253Z","iopub.status.idle":"2020-08-30T16:56:06.055536Z","shell.execute_reply":"2020-08-30T16:56:06.056115Z"},"papermill":{"duration":0.021999,"end_time":"2020-08-30T16:56:06.056226","exception":false,"start_time":"2020-08-30T16:56:06.034227","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"def prediction(img):\n    a = rgb2lab(img.permute(1, 2, 0))\n    a = torch.tensor(a[:,:,0]).type(torch.FloatTensor)\n    a = a.unsqueeze(0)\n    a = a.unsqueeze(0)\n    xb = to_device(a, device)\n    ab_img = model(xb)\n    xb = xb.squeeze(0)\n    ab_img = ab_img.squeeze(0)\n    return to_rgb(xb.detach().cpu(), ab_img.detach().cpu())","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.013341,"end_time":"2020-08-30T16:56:06.082079","exception":false,"start_time":"2020-08-30T16:56:06.068738","status":"completed"},"tags":[]},"cell_type":"markdown","source":"Testing the model on different black and white images and plotting the result."},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:56:06.111184Z","iopub.status.busy":"2020-08-30T16:56:06.110610Z","iopub.status.idle":"2020-08-30T16:56:06.116922Z","shell.execute_reply":"2020-08-30T16:56:06.116354Z"},"papermill":{"duration":0.022138,"end_time":"2020-08-30T16:56:06.117012","exception":false,"start_time":"2020-08-30T16:56:06.094874","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"import glob\nfrom PIL import Image\n\nimages = glob.glob(\"../input/grayscale-landscape/download (7).jpg\")","execution_count":null,"outputs":[]},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:56:06.146558Z","iopub.status.busy":"2020-08-30T16:56:06.145880Z","iopub.status.idle":"2020-08-30T16:56:06.198616Z","shell.execute_reply":"2020-08-30T16:56:06.198105Z"},"papermill":{"duration":0.069373,"end_time":"2020-08-30T16:56:06.198708","exception":false,"start_time":"2020-08-30T16:56:06.129335","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"img = images[0]\nimage = Image.open(img)\ntrans = T.Compose([T.Resize((256, 256)),T.ToTensor()])\nimg = trans(image)","execution_count":null,"outputs":[]},{"metadata":{"execution":{"iopub.execute_input":"2020-08-30T16:56:06.229801Z","iopub.status.busy":"2020-08-30T16:56:06.229093Z","iopub.status.idle":"2020-08-30T16:56:07.280012Z","shell.execute_reply":"2020-08-30T16:56:07.280744Z"},"papermill":{"duration":1.069003,"end_time":"2020-08-30T16:56:07.280914","exception":false,"start_time":"2020-08-30T16:56:06.211911","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"f, arr = plt.subplots(1, 2, sharey=True)\narr[0].imshow(img.permute(1, 2, 0))\narr[1].imshow(prediction(img))","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.024852,"end_time":"2020-08-30T16:56:07.330037","exception":false,"start_time":"2020-08-30T16:56:07.305185","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.022592,"end_time":"2020-08-30T16:56:07.374180","exception":false,"start_time":"2020-08-30T16:56:07.351588","status":"completed"},"tags":[],"trusted":false},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}