initial commit

Pipfile

@@ -0,0 +1,11 @@
+[[source]]
+url = "https://pypi.org/simple"
+verify_ssl = true
+name = "pypi"
+
+[packages]
+
+[dev-packages]
+
+[requires]
+python_version = "3.10"

main.py

@@ -0,0 +1,65 @@
+import warnings
+warnings.filterwarnings("ignore")
+
+import os
+import sys
+import glob
+import time
+import numpy as np
+from PIL import Image
+from pathlib import Path
+from tqdm.notebook import tqdm
+import matplotlib.pyplot as plt
+from skimage.color import rgb2lab, lab2rgb
+
+import torch
+from torch import nn, optim
+from torchvision import transforms
+from torchvision.utils import make_grid
+from torch.utils.data import Dataset, DataLoader
+
+
+from utility import *
+from model import *
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model_path = "model/ImageColorizationModel.pth"
+
+
+model = None
+if not os.path.exists(model_path) :
+    print("Model not find")
+    download_from_drive()
+    print("Model Downloaded")
+else:
+    model = load_model(model_class=MainModel , file_path=model_path)
+    print("Model Loaded")
+
+def predict_and_return_image(image):
+    data = create_lab_tensors(image)
+    model.net_G.eval()
+    with torch.no_grad():
+        model.setup_input(data)
+        model.forward()
+    fake_color = model.fake_color.detach()
+    L = model.L
+    fake_imgs = lab_to_rgb(L, fake_color)
+    return fake_imgs[0]
+
+
+
+
+
+title = "Black&White to Color image"
+description = "Transforming Black & White Image in to colored image. Upload a black and white image to see it colorized by our deep learning model."
+
+gr.Interface(
+    fn=predict_and_return_image,
+    title=title,
+    description=description,
+    inputs=[gr.Image(label="Gray Scale Image")],
+    outputs=[
+        gr.Image(label="Predicted Colored Image")
+    ],
+).launch(share=True, debug=True)

model/Discriminator.py

@@ -0,0 +1,37 @@
+import os
+import glob
+import time
+import numpy as np
+from PIL import Image
+from pathlib import Path
+from tqdm.notebook import tqdm
+import matplotlib.pyplot as plt
+from skimage.color import rgb2lab, lab2rgb
+
+import torch
+from torch import nn, optim
+from torchvision import transforms
+from torchvision.utils import make_grid
+from torch.utils.data import Dataset, DataLoader
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+class PatchDiscriminator(nn.Module):
+    def __init__(self, input_c, num_filters=64, n_down=3):
+        super().__init__()
+        model = [self.get_layers(input_c, num_filters, norm=False)]
+        model += [self.get_layers(num_filters * 2 ** i, num_filters * 2 ** (i + 1), s=1 if i == (n_down-1) else 2) 
+                          for i in range(n_down)] # the 'if' statement is taking care of not using
+                                                  # stride of 2 for the last block in this loop
+        model += [self.get_layers(num_filters * 2 ** n_down, 1, s=1, norm=False, act=False)] # Make sure to not use normalization or
+                                                                                             # activation for the last layer of the model
+        self.model = nn.Sequential(*model)                                                   
+        
+    def get_layers(self, ni, nf, k=4, s=2, p=1, norm=True, act=True): # when needing to make some repeatitive blocks of layers,
+        layers = [nn.Conv2d(ni, nf, k, s, p, bias=not norm)]          # it's always helpful to make a separate method for that purpose
+        if norm: layers += [nn.BatchNorm2d(nf)]
+        if act: layers += [nn.LeakyReLU(0.2, True)]
+        return nn.Sequential(*layers)
+    
+    def forward(self, x):
+        return self.model(x)

model/Generator.py

@@ -0,0 +1,73 @@
+import os
+import glob
+import time
+import numpy as np
+from PIL import Image
+from pathlib import Path
+from tqdm.notebook import tqdm
+import matplotlib.pyplot as plt
+from skimage.color import rgb2lab, lab2rgb
+
+import torch
+from torch import nn, optim
+from torchvision import transforms
+from torchvision.utils import make_grid
+from torch.utils.data import Dataset, DataLoader
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+class UnetBlock(nn.Module):
+    def __init__(self, nf, ni, submodule=None, input_c=None, dropout=False,
+                 innermost=False, outermost=False):
+        super().__init__()
+        self.outermost = outermost
+        if input_c is None: input_c = nf
+        downconv = nn.Conv2d(input_c, ni, kernel_size=4,
+                             stride=2, padding=1, bias=False)
+        downrelu = nn.LeakyReLU(0.2, True)
+        downnorm = nn.BatchNorm2d(ni)
+        uprelu = nn.ReLU(True)
+        upnorm = nn.BatchNorm2d(nf)
+        
+        if outermost:
+            upconv = nn.ConvTranspose2d(ni * 2, nf, kernel_size=4,
+                                        stride=2, padding=1)
+            down = [downconv]
+            up = [uprelu, upconv, nn.Tanh()]
+            model = down + [submodule] + up
+        elif innermost:
+            upconv = nn.ConvTranspose2d(ni, nf, kernel_size=4,
+                                        stride=2, padding=1, bias=False)
+            down = [downrelu, downconv]
+            up = [uprelu, upconv, upnorm]
+            model = down + up
+        else:
+            upconv = nn.ConvTranspose2d(ni * 2, nf, kernel_size=4,
+                                        stride=2, padding=1, bias=False)
+            down = [downrelu, downconv, downnorm]
+            up = [uprelu, upconv, upnorm]
+            if dropout: up += [nn.Dropout(0.5)]
+            model = down + [submodule] + up
+        self.model = nn.Sequential(*model)
+    
+    def forward(self, x):
+        if self.outermost:
+            return self.model(x)
+        else:
+            return torch.cat([x, self.model(x)], 1)
+
+
+class Unet(nn.Module):
+    def __init__(self, input_c=1, output_c=2, n_down=8, num_filters=64):
+        super().__init__()
+        unet_block = UnetBlock(num_filters * 8, num_filters * 8, innermost=True)
+        for _ in range(n_down - 5):
+            unet_block = UnetBlock(num_filters * 8, num_filters * 8, submodule=unet_block, dropout=True)
+        out_filters = num_filters * 8
+        for _ in range(3):
+            unet_block = UnetBlock(out_filters // 2, out_filters, submodule=unet_block)
+            out_filters //= 2
+        self.model = UnetBlock(output_c, out_filters, input_c=input_c, submodule=unet_block, outermost=True)
+    
+    def forward(self, x):
+        return self.model(x)

model/__init__.py

@@ -0,0 +1,87 @@
+import os
+import glob
+import time
+import numpy as np
+from PIL import Image
+from pathlib import Path
+from tqdm.notebook import tqdm
+import matplotlib.pyplot as plt
+from skimage.color import rgb2lab, lab2rgb
+
+import torch
+from torch import nn, optim
+from torchvision import transforms
+from torchvision.utils import make_grid
+from torch.utils.data import Dataset, DataLoader
+
+from .Generator import UnetBlock , Unet
+from .Discriminator import PatchDiscriminator
+from .weights import init_weights
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def init_model(model, device):
+    model = model.to(device)
+    model = init_weights(model)
+    return model
+
+class MainModel(nn.Module):
+    def __init__(self, net_G=None, lr_G=2e-4, lr_D=2e-4, 
+                 beta1=0.5, beta2=0.999, lambda_L1=100.):
+        super().__init__()
+        
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.lambda_L1 = lambda_L1
+        
+        if net_G is None:
+            self.net_G = init_model(Unet(input_c=1, output_c=2, n_down=8, num_filters=64), self.device)
+        else:
+            self.net_G = net_G.to(self.device)
+        self.net_D = init_model(PatchDiscriminator(input_c=3, n_down=3, num_filters=64), self.device)
+        self.GANcriterion = GANLoss(gan_mode='vanilla').to(self.device)
+        self.L1criterion = nn.L1Loss()
+        self.opt_G = optim.Adam(self.net_G.parameters(), lr=lr_G, betas=(beta1, beta2))
+        self.opt_D = optim.Adam(self.net_D.parameters(), lr=lr_D, betas=(beta1, beta2))
+    
+    def set_requires_grad(self, model, requires_grad=True):
+        for p in model.parameters():
+            p.requires_grad = requires_grad
+        
+    def setup_input(self, data):
+        self.L = data['L'].to(self.device)
+        self.ab = data['ab'].to(self.device)
+        
+    def forward(self):
+        self.fake_color = self.net_G(self.L)
+    
+    def backward_D(self):
+        fake_image = torch.cat([self.L, self.fake_color], dim=1)
+        fake_preds = self.net_D(fake_image.detach())
+        self.loss_D_fake = self.GANcriterion(fake_preds, False)
+        real_image = torch.cat([self.L, self.ab], dim=1)
+        real_preds = self.net_D(real_image)
+        self.loss_D_real = self.GANcriterion(real_preds, True)
+        self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5
+        self.loss_D.backward()
+    
+    def backward_G(self):
+        fake_image = torch.cat([self.L, self.fake_color], dim=1)
+        fake_preds = self.net_D(fake_image)
+        self.loss_G_GAN = self.GANcriterion(fake_preds, True)
+        self.loss_G_L1 = self.L1criterion(self.fake_color, self.ab) * self.lambda_L1
+        self.loss_G = self.loss_G_GAN + self.loss_G_L1
+        self.loss_G.backward()
+    
+    def optimize(self):
+        self.forward()
+        self.net_D.train()
+        self.set_requires_grad(self.net_D, True)
+        self.opt_D.zero_grad()
+        self.backward_D()
+        self.opt_D.step()
+        
+        self.net_G.train()
+        self.set_requires_grad(self.net_D, False)
+        self.opt_G.zero_grad()
+        self.backward_G()
+        self.opt_G.step()

model/weights.py

@@ -0,0 +1,39 @@
+import os
+import glob
+import time
+import numpy as np
+from PIL import Image
+from pathlib import Path
+from tqdm.notebook import tqdm
+import matplotlib.pyplot as plt
+from skimage.color import rgb2lab, lab2rgb
+
+import torch
+from torch import nn, optim
+from torchvision import transforms
+from torchvision.utils import make_grid
+from torch.utils.data import Dataset, DataLoader
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def init_weights(net, init='norm', gain=0.02):
+    
+    def init_func(m):
+        classname = m.__class__.__name__
+        if hasattr(m, 'weight') and 'Conv' in classname:
+            if init == 'norm':
+                nn.init.normal_(m.weight.data, mean=0.0, std=gain)
+            elif init == 'xavier':
+                nn.init.xavier_normal_(m.weight.data, gain=gain)
+            elif init == 'kaiming':
+                nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+            
+            if hasattr(m, 'bias') and m.bias is not None:
+                nn.init.constant_(m.bias.data, 0.0)
+        elif 'BatchNorm2d' in classname:
+            nn.init.normal_(m.weight.data, 1., gain)
+            nn.init.constant_(m.bias.data, 0.)
+            
+    net.apply(init_func)
+    print(f"model initialized with {init} initialization")
+    return net

utility/__init__.py

@@ -0,0 +1 @@
+from .helper import *

utility/helper.py

@@ -0,0 +1,181 @@
+import os
+import glob
+import time
+import numpy as np
+from PIL import Image
+from pathlib import Path
+from tqdm.notebook import tqdm
+import matplotlib.pyplot as plt
+from skimage.color import rgb2lab, lab2rgb
+
+import torch
+from torch import nn, optim
+from torchvision import transforms
+from torchvision.utils import make_grid
+from torch.utils.data import Dataset, DataLoader
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+import requests
+import gdown
+
+def download_from_drive():
+    url = "https://drive.google.com/uc?id=1EhuMET76c02VFyRW8Pie7BwNCDHmQiad"
+    try:
+        output = "model/ImageColorizationModel.pth"
+        gdown.download(url, output, quiet=False)
+        return True
+    except:
+        print("Error Occured in Downloading model from Gdrive")
+        return False
+
+
+class AverageMeter:
+    def __init__(self):
+        self.reset()
+        
+    def reset(self):
+        self.count, self.avg, self.sum = [0.] * 3
+    
+    def update(self, val, count=1):
+        self.count += count
+        self.sum += count * val
+        self.avg = self.sum / self.count
+
+def create_loss_meters():
+    loss_D_fake = AverageMeter()
+    loss_D_real = AverageMeter()
+    loss_D = AverageMeter()
+    loss_G_GAN = AverageMeter()
+    loss_G_L1 = AverageMeter()
+    loss_G = AverageMeter()
+    
+    return {'loss_D_fake': loss_D_fake,
+            'loss_D_real': loss_D_real,
+            'loss_D': loss_D,
+            'loss_G_GAN': loss_G_GAN,
+            'loss_G_L1': loss_G_L1,
+            'loss_G': loss_G}
+
+def update_losses(model, loss_meter_dict, count):
+    for loss_name, loss_meter in loss_meter_dict.items():
+        loss = getattr(model, loss_name)
+        loss_meter.update(loss.item(), count=count)
+
+def lab_to_rgb(L, ab):
+    """
+    Takes a batch of images
+    """
+    
+    L = (L + 1.) * 50.
+    ab = ab * 110.
+    Lab = torch.cat([L, ab], dim=1).permute(0, 2, 3, 1).cpu().numpy()
+    rgb_imgs = []
+    for img in Lab:
+        img_rgb = lab2rgb(img)
+        rgb_imgs.append(img_rgb)
+    return np.stack(rgb_imgs, axis=0)
+    
+def visualize(model, data, save=True):
+    model.net_G.eval()
+    with torch.no_grad():
+        model.setup_input(data)
+        model.forward()
+    model.net_G.train()
+    fake_color = model.fake_color.detach()
+    real_color = model.ab
+    L = model.L
+    fake_imgs = lab_to_rgb(L, fake_color)
+    real_imgs = lab_to_rgb(L, real_color)
+    fig = plt.figure(figsize=(15, 8))
+    for i in range(5):
+        ax = plt.subplot(3, 5, i + 1)
+        ax.imshow(L[i][0].cpu(), cmap='gray')
+        ax.axis("off")
+        ax = plt.subplot(3, 5, i + 1 + 5)
+        ax.imshow(fake_imgs[i])
+        ax.axis("off")
+        ax = plt.subplot(3, 5, i + 1 + 10)
+        ax.imshow(real_imgs[i])
+        ax.axis("off")
+    plt.show()
+    if save:
+        fig.savefig(f"colorization_{time.time()}.png")
+        
+def log_results(loss_meter_dict):
+    for loss_name, loss_meter in loss_meter_dict.items():
+        print(f"{loss_name}: {loss_meter.avg:.5f}")
+        
+def create_lab_tensors(image):
+    """
+    This function receives an image path or a direct image input and creates a dictionary of L and ab tensors.
+    Args:
+    - image: either a path to the image file or a direct image input.
+    Returns:
+    - lab_dict: dictionary containing the L and ab tensors.
+    """
+    if isinstance(image, str):
+        # Open the image and convert it to RGB format
+        img = Image.open(image).convert('RGB')
+    else:
+        img = image.convert('RGB')
+        
+    custom_transforms = transforms.Compose([
+                transforms.Resize((SIZE, SIZE),  Image.BICUBIC),
+                transforms.RandomHorizontalFlip(), # A little data augmentation!
+            ])
+    img = custom_transforms(img)
+    img = np.array(img)
+    img_lab = rgb2lab(img).astype("float32") # Converting RGB to L*a*b
+    img_lab = transforms.ToTensor()(img_lab)
+    L = img_lab[[0], ...] / 50. - 1. # Between -1 and 1
+    L = L.unsqueeze(0) 
+    ab = img_lab[[1, 2], ...] / 110. # Between -1 and 1
+    return {'L': L, 'ab': ab}
+
+
+def predict_and_visualize_single_image(model, data, save=True):
+    model.net_G.eval()
+    with torch.no_grad():
+        model.setup_input(data)
+        model.forward()
+    fake_color = model.fake_color.detach()
+    L = model.L
+    fake_imgs = lab_to_rgb(L, fake_color)
+    fig, axs = plt.subplots(1, 2, figsize=(8, 4))
+    axs[0].imshow(L[0][0].cpu(), cmap='gray')
+    axs[0].set_title("Grey Image")
+    axs[0].axis('off')
+
+    axs[1].imshow(fake_imgs[0])
+    axs[1].set_title("Colored Image")
+    axs[1].axis('off')
+    plt.show()
+    if save:
+        fig.savefig(f"colorization_{time.time()}.png")
+        
+def predict_color(model , image , save=False):
+    """
+    This function receives an image path or a direct image input and creates a dictionary of L and ab tensors.
+    Args:
+    - model : Pytorch Gray Scale to Colorization Model
+    - image: either a path to the image file or a direct image input.
+    """
+    data = create_lab_tensors(image)
+    predict_and_visualize_single_image(model, data, save)
+
+
+def load_model(model_class, file_path):
+    """
+    Load PyTorch model from file.
+    
+    Args:
+        model_class (torch.nn.Module): PyTorch model class to load.
+        file_path (str): File path to load the model from.
+    
+    Returns:
+        model (torch.nn.Module): Loaded PyTorch model.
+    """
+    model = model_class()
+    model.load_state_dict(torch.load(file_path))
+    return model