!pip install git+https://gitlab.aicrowd.com/yoogottamk/aicrowd-cli.git
from aicrowd.magic import *
# Get your API key from https://www.aicrowd.com/participants/me
%aicrowd login --api-key <key>

%aicrowd dataset download --challenge imgcol *v2.zip

!rm -rf data
!mkdir data

!unzip train_black_white_images-v2.zip  -d data/
!unzip train_color_images-v2.zip -d data/
!unzip validation_black_white_images-v2.zip  -d data/
!unzip validation_color_images-v2.zip  -d data/
!unzip test_black_white_images-v2.zip -d data/

import torch
import os
import cv2
import albumentations as A
from torch.utils.data import DataLoader, Dataset, Subset
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
import numpy as np
import matplotlib.pyplot as plt
import shutil
from tqdm.notebook import tqdm
from PIL import Image

class CreateDataset(Dataset):
  
  def __init__(self, gray_path, color_path=None):
    self.gray_path = gray_path
    self.color_path = color_path

    self.gray_imgs_paths = os.listdir(self.gray_path)

  def __getitem__(self, idx):
    
    if self.color_path:

      gray_img = cv2.imread(f"{self.gray_path}/{self.gray_imgs_paths[idx]}", 0)
      color_img = cv2.imread(f"{self.color_path}/{self.gray_imgs_paths[idx]}")

      # gray_img = cv2.resize(gray_img, (512, 512))
      # color_img = cv2.resize(color_img, (512, 512))

      gray_img = gray_img[np.newaxis, :, :]

      gray_img = torch.tensor(gray_img)
      color_img = torch.tensor(color_img)

      return gray_img, color_img.permute(2, 0, 1)

    else:
      
      gray_img = cv2.imread(f"{self.gray_path}/{self.gray_imgs_paths[idx]}", 0)

      gray_img = gray_img[np.newaxis, :, :]

      gray_img = torch.tensor(gray_img)
      img_id = self.gray_imgs_paths[idx]

      return gray_img, img_id
  
  
  def __len__(self):

    return len(self.gray_imgs_paths)

train_class = CreateDataset("data/train_black_white_images-v2", "data/train_color_images-v2")
train_loader = torch.utils.data.DataLoader(train_class, batch_size=4, shuffle=True)

test_class = CreateDataset("data/test_black_white_images-v2")
test_loader = torch.utils.data.DataLoader(test_class, batch_size=4, shuffle=True)

class ColorizationNet(nn.Module0bcde28da0f5f8dbbff8e460627c3c48 ):

  """
  From https://lukemelas.github.io/image-colorization.html
  """

  def __init__(self, input_size=512):
    super(ColorizationNet, self).__init__()
    MIDLEVEL_FEATURE_SIZE = 128

 
    resnet = models.resnet18(num_classes=365) 
    resnet.conv1.weight = nn.Parameter(resnet.conv1.weight.sum(dim=1).unsqueeze(1)) 
    self.midlevel_resnet = nn.Sequential(*list(resnet.children())[0:6])

    self.upsample = nn.Sequential(     
      nn.Conv2d(MIDLEVEL_FEATURE_SIZE, 128, kernel_size=3, stride=1, padding=1),
      nn.BatchNorm2d(128),
      nn.ReLU(),
      nn.Upsample(scale_factor=2),
      nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1),
      nn.BatchNorm2d(64),
      nn.ReLU(),
      nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
      nn.BatchNorm2d(64),
      nn.ReLU(),
      nn.Upsample(scale_factor=2),
      nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1),
      nn.BatchNorm2d(32),
      nn.ReLU(),
      nn.Conv2d(32, 3, kernel_size=3, stride=1, padding=1),
      nn.Upsample(scale_factor=2)
    )

  def forward(self, input):

    midlevel_features = self.midlevel_resnet(input)
    output = self.upsample(midlevel_features)
    return output

model = ColorizationNet()

# Putting sample image and getting predictions

sample_img = torch.tensor(np.random.rand(1, 1, 512, 512)).float()
plt.imshow(sample_img[0][0])
prediction = model(sample_img)

plt.imshow(prediction.detach().numpy()[0].transpose(2, 1, 0))

# Setting up optimizer & loss

optimizer = torch.optim.Adam(model.parameters(), lr=1e-2, weight_decay=0.0)
criterion = nn.MSELoss()

use_gpu = torch.cuda.is_available()

if use_gpu: 
  criterion = criterion.cuda()
  model = model.cuda()

model.train()

for i, (input_gray, target) in enumerate(train_loader):

    if use_gpu: input_gray, target = input_gray.cuda().float(), target.cuda().float()

    prediction = model(input_gray) 
    loss = criterion(prediction.float(), target) 
    #losses.update(loss.item(), input_gray.size(0))

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if i % 50 == 0:
      print("loss", loss)

%rm -rf submission

test_imgs_paths = os.listdir("data/test_black_white_images-v2")
os.makedirs("test_color_images", exist_ok=True)

model.eval()

with torch.no_grad():

  for i, (input_gray, img_id) in enumerate(tqdm(test_loader)):

    predictions = model(input_gray.cuda().float())

    for n, prediction in enumerate(predictions):
      
      prediction = prediction.detach().cpu().numpy().transpose(2, 1, 0)

      cv2.imwrite(f"test_color_images/{img_id[n]}", prediction)

!zip -r "first_submission.zip" "test_color_images/"

try:
    from google.colab import files
    files.download('first_submission.zip') 
except:
    print("Option Only avilable in Google Colab")