{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"VERSION = \"20200220\" #@param [\"20200220\",\"nightly\", \"xrt==1.15.0\"]\n!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n!python pytorch-xla-env-setup.py --version $VERSION","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#!pip install -U git+https://github.com/albumentations-team/albumentations\nimport torch \nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.functional as F\nfrom torch.autograd import Variable\nimport numpy as np\nimport os \nimport pathlib\nfrom torch.utils.data import Dataset,DataLoader\nimport time\nimport matplotlib.pyplot as plt\nfrom torchvision import transforms\nimport cv2\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport datetime\nimport sys\n#import torch_xla\n#import torch_xla.core.xla_model as xm","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"image_paths_flickr = [x for x in pathlib.Path('../input/flickr-image-dataset').rglob('*.jpg')]\nimage_paths_landmark = [x for x in pathlib.Path('../input/landmark-recognition-2020/train').rglob('*.jpg')]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"f\"{len(image_paths_flickr)},{len(image_paths_landmark)}\"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"all_image_paths = image_paths_flickr + image_paths_landmark\nall_image_paths = [str(i) for i in all_image_paths]\nf\"{len(all_image_paths)} {all_image_paths[0]}\"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def show_image(arr):\n    arr = np.moveaxis(arr,0,-1)\n    plt.imshow(arr)\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"img = cv2.imread(all_image_paths[-1])\nimg = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)\nimg = np.moveaxis(img,-1,0) # Channel first\nprint(img.shape)\nshow_image(img)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"gimg = cv2.cvtColor(np.moveaxis(img,0,-1),cv2.COLOR_RGB2GRAY)\ngimg = np.stack((gimg,gimg,gimg),axis=0)\n\nprint(gimg.shape)\nshow_image(gimg)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"transforms = A.Compose([\n    A.Resize(256,256),\n    A.Normalize(),\n    A.pytorch.transforms.ToTensorV2()\n            ])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ColorDataset(Dataset):\n    def __init__(self,img_list,transform):\n        self.img_list = img_list\n        self.transform = transform\n    \n    def __len__(self):\n        return len(self.img_list)\n    \n    def __getitem__(self,idx):\n        if torch.is_tensor(idx): idx = idx.tolist()\n        img = cv2.imread(all_image_paths[idx])\n        img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)\n        gimg = cv2.cvtColor(img,cv2.COLOR_RGB2GRAY)\n        gimg = np.stack((gimg,gimg,gimg),axis=-1)\n\n        img = self.transform(image = img)[\"image\"]\n        gimg = self.transform(image=gimg)[\"image\"]\n\n        return img,gimg\n    \n            ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"Cds = ColorDataset(all_image_paths,transforms)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"Cds[0][0].shape","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def weights_init_normal(m):\n    classname = m.__class__.__name__\n    if classname.find(\"Conv\") != -1:\n        torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n    elif classname.find(\"BatchNorm2d\") != -1:\n        torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n        torch.nn.init.constant_(m.bias.data, 0.0)\n\n\n##############################\n#           U-NET\n##############################\n\n\nclass UNetDown(nn.Module):\n    def __init__(self, in_size, out_size, normalize=True, dropout=0.0):\n        super(UNetDown, self).__init__()\n        layers = [nn.Conv2d(in_size, out_size, 4, 2, 1, bias=False)]\n        if normalize:\n            layers.append(nn.InstanceNorm2d(out_size))\n        layers.append(nn.LeakyReLU(0.2))\n        if dropout:\n            layers.append(nn.Dropout(dropout))\n        self.model = nn.Sequential(*layers)\n\n    def forward(self, x):\n        return self.model(x)\n\n\nclass UNetUp(nn.Module):\n    def __init__(self, in_size, out_size, dropout=0.0):\n        super(UNetUp, self).__init__()\n        layers = [\n            nn.ConvTranspose2d(in_size, out_size, 4, 2, 1, bias=False),\n            nn.InstanceNorm2d(out_size),\n            nn.ReLU(inplace=True),\n        ]\n        if dropout:\n            layers.append(nn.Dropout(dropout))\n\n        self.model = nn.Sequential(*layers)\n\n    def forward(self, x, skip_input):\n        x = self.model(x)\n        x = torch.cat((x, skip_input), 1)\n\n        return x\n\n\nclass generator(nn.Module):\n    def __init__(self, in_channels=3, out_channels=3):\n        super(generator, self).__init__()\n\n        self.down1 = UNetDown(in_channels, 64, normalize=False)\n        self.down2 = UNetDown(64, 128)\n        self.down3 = UNetDown(128, 256)\n        self.down4 = UNetDown(256, 512, dropout=0.5)\n        self.down5 = UNetDown(512, 512, dropout=0.5)\n        self.down6 = UNetDown(512, 512, dropout=0.5)\n        self.down7 = UNetDown(512, 512, dropout=0.5)\n        self.down8 = UNetDown(512, 512, normalize=False, dropout=0.5)\n\n        self.up1 = UNetUp(512, 512, dropout=0.5)\n        self.up2 = UNetUp(1024, 512, dropout=0.5)\n        self.up3 = UNetUp(1024, 512, dropout=0.5)\n        self.up4 = UNetUp(1024, 512, dropout=0.5)\n        self.up5 = UNetUp(1024, 256)\n        self.up6 = UNetUp(512, 128)\n        self.up7 = UNetUp(256, 64)\n\n        self.final = nn.Sequential(\n            nn.Upsample(scale_factor=2),\n            nn.ZeroPad2d((1, 0, 1, 0)),\n            nn.Conv2d(128, out_channels, 4, padding=1),\n            nn.Tanh(),\n        )\n\n    def forward(self, x):\n        # U-Net generator with skip connections from encoder to decoder\n        d1 = self.down1(x)\n        d2 = self.down2(d1)\n        d3 = self.down3(d2)\n        d4 = self.down4(d3)\n        d5 = self.down5(d4)\n        d6 = self.down6(d5)\n        d7 = self.down7(d6)\n        d8 = self.down8(d7)\n        u1 = self.up1(d8, d7)\n        u2 = self.up2(u1, d6)\n        u3 = self.up3(u2, d5)\n        u4 = self.up4(u3, d4)\n        u5 = self.up5(u4, d3)\n        u6 = self.up6(u5, d2)\n        u7 = self.up7(u6, d1)\n\n        return self.final(u7)\n\n\n##############################\n#        Discriminator\n##############################\n\n\nclass discriminator(nn.Module):\n    def __init__(self, in_channels=3):\n        super(discriminator, self).__init__()\n\n        def discriminator_block(in_filters, out_filters, normalization=True):\n            \"\"\"Returns downsampling layers of each discriminator block\"\"\"\n            layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]\n            if normalization:\n                layers.append(nn.InstanceNorm2d(out_filters))\n            layers.append(nn.LeakyReLU(0.2, inplace=True))\n            return layers\n\n        self.model = nn.Sequential(\n            *discriminator_block(in_channels * 2, 64, normalization=False),\n            *discriminator_block(64, 128),\n            *discriminator_block(128, 256),\n            *discriminator_block(256, 512),\n            nn.ZeroPad2d((1, 0, 1, 0)),\n            nn.Conv2d(512, 1, 4, padding=1, bias=False)\n        )\n\n    def forward(self, img_A, img_B):\n        # Concatenate image and condition image by channels to produce input\n        img_input = torch.cat((img_A, img_B), 1)\n        return self.model(img_input)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"netG = generator()\nnetD = discriminator()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"lr = 0.001\nbatch_size = 32\nn_epochs = 1\nimg_height = 256\nimg_width = 256\noptim_G = torch.optim.Adam(netG.parameters(),lr=lr)\noptim_D = torch.optim.Adam(netD.parameters(),lr=lr)\ndevice = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\ncriterion_GAN = torch.nn.MSELoss()\ncriterion_pixelwise = torch.nn.L1Loss()\npatch = (1, img_height // 2 ** 4, img_width // 2 ** 4)#loss of image disc patch gan\nlambda_pixel = 100\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"netG = netG.to(device)\nnetD = netD.to(device)\ncriterion_GAN.to(device)\ncriterion_pixelwise.to(device)\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_hist = {}\ntrain_hist['D_losses'] = []\ntrain_hist['G_losses'] = []\ntrain_hist['per_epoch_ptimes'] = []\ntrain_hist['total_ptime'] = []","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"imdl = DataLoader(Cds,batch_size = batch_size,shuffle=True)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def sample_images(batches_done):\n    \"\"\"Saves a generated sample from the validation set\"\"\"\n    imgs = next(iter(imdl))\n    real_A = Variable(imgs[1].type(Tensor))\n    real_B = Variable(imgs[0].type(Tensor))\n    fake_B = netG(real_A)\n    img_sample = torch.cat((real_A.data, fake_B.data, real_B.data), -2).cpu().numpy().astype(np.float32)\n    img_sample -=img_sample.min()\n    img_sample/=img_sample.max()\n    img_sample = img_sample.transpose(0,2,3,1)\n    plt.figure(figsize=[10,20])\n    for row in range(3):\n        plt.subplot(1,3,row+1)\n        plt.imshow(img_sample[row])\n    plt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"Tensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"prev_time = time.time()\nsample_interval = 500\ncheckpoint_interval = 338\nfor epoch in range(n_epochs):\n    for i, batch in enumerate(imdl):\n\n        # Model inputs\n        real_A= Variable(batch[1].type(Tensor))#black\n        real_B= Variable(batch[0].type(Tensor))#color\n\n        # Adversarial ground truths\n        valid = Variable(Tensor(np.ones((real_A.size(0), *patch))), requires_grad=False)#doubt here(loss of umage disc patch gan)\n        fake = Variable(Tensor(np.zeros((real_A.size(0), *patch))), requires_grad=False)\n\n        # ------------------\n        #  Train Generators\n        # ------------------\n\n        optim_G.zero_grad()\n\n        # GAN loss\n        fake_B = netG(real_A)\n        pred_fake = netD(fake_B, real_A)\n        loss_GAN = criterion_GAN(pred_fake, valid)\n        # Pixel-wise loss\n        loss_pixel = criterion_pixelwise(fake_B, real_B)\n\n        # Total loss\n        loss_G = loss_GAN + lambda_pixel * loss_pixel\n\n        loss_G.backward()\n\n        optim_G.step()\n\n        # ---------------------\n        #  Train Discriminator\n        # ---------------------\n\n        optim_D.zero_grad()\n\n        # Real loss\n        pred_real = netD(real_B, real_A)\n        loss_real = criterion_GAN(pred_real, valid)\n\n        # Fake loss\n        pred_fake = netD(fake_B.detach(), real_A)\n        loss_fake = criterion_GAN(pred_fake, fake)\n\n        # Total loss\n        loss_D = 0.5 * (loss_real + loss_fake)\n\n        loss_D.backward()\n        optim_D.step()\n\n        # --------------\n        #  Log Progress\n        # --------------\n\n        # Determine approximate time left\n        batches_done = epoch * len(imdl) + i\n        batches_left = n_epochs * len(imdl) - batches_done\n        time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))\n        prev_time = time.time()\n\n        # Print log\n        sys.stdout.write(\n            \"\\r[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f, pixel: %f, adv: %f] ETA: %s\"\n            % (\n                epoch,\n                n_epochs,\n                i,\n                len(imdl),\n                loss_D.item(),\n                loss_G.item(),\n                loss_pixel.item(),\n                loss_GAN.item(),\n                time_left,\n            )\n        )\n\n        # If at sample interval save image\n        if batches_done % sample_interval == 0:\n            sample_images(batches_done)\n        if batches_done % 5000 == 0:\n            torch.save(netG.state_dict(), f'colourizeGen{n_epochs}.pt')\n            torch.save(netD.state_dict(), f'colourizeDis{n_epochs}.pt')\n            print(\"model saved \")\n            \n            \n\n  #  if checkpoint_interval != -1 and epoch % checkpoint_interval == 0:\n   #     # Save model checkpoints\n     #   torch.save(netG.state_dict(), f'colourizeGen{epoch+20}.pth')\n     #   torch.save(netD.state_dict(), f'colourizeDis{epoch+20}.pth')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a href=\"./colourizeGen1.pt\"> Download File </a>","metadata":{}},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"markdown","source":"## ","metadata":{}},{"cell_type":"code","source":"torch.save(netG.state_dict(), f'colourizeGen{n_epoch}.pt')\ntorch.save(netD.state_dict(), f'colourizeDis{n_epoch}.pt')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}