add everything

app.py

@@ -0,0 +1,49 @@
+import gradio as gr
+from PIL import Image
+import os
+import numpy as np
+# from outpaint import outpainting
+# from model import colorazation, UNETmodel, utils1
+# from model import inference, model
+# from model import colorazation, deeplabmodel, utils 
+from model import MainModel
+import inference as inf
+
+
+# pretrained model
+def colorize_image(image):
+     # Load the model
+    # file_path = 'ImageColorizationModel10.pth'
+    file_path = r'model\model_final.pth'
+    model_2 = inf.load_model(model_class=MainModel, file_path=file_path)
+    output_img = inf.predict_color(model_2, image=image)
+    return output_img
+
+
+
+# pretrained model
+colorization_interface = gr.Interface(
+    colorize_image,
+    gr.Image(type="pil", label="Input Image"),
+    [gr.Image(type="pil", label="Output Image")],
+    title="Image Colorization",
+    description="Upload an image to perform colorization.",
+
+)
+
+# deeplab model
+# depinterface = gr.Interface(
+#     depColorize_image,
+#     gr.Image(type="pil", label="Input Image"),
+#     [gr.Image(type="pil", label="Output Image")],
+#     title="Image Colorization",
+#     description="Upload an image to perform colorization.",
+
+# )
+
+# scratch mod
+
+# Launch the interface
+# interface.launch(share=True)
+with gr.TabbedInterface([ colorization_interface ], ["Colorization_pretrain_unet"]) as tabs:
+    tabs.launch(share=True)

inference.py

@@ -0,0 +1,129 @@
+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
+
+
+def init_model(model, device):
+    model = model.to(device)
+    model = init_weights(model)
+    return model
+
+
+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
+
+from fastai.vision.learner import create_body
+from torchvision.models.resnet import resnet18
+from fastai.vision.models.unet import DynamicUnet
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+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 build_res_unet(n_input=1, n_output=2, size=256):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    body = create_body(resnet18(), pretrained=True, n_in=n_input, cut=-2)
+    net_G = DynamicUnet(body, n_output, (size, size)).to(device)
+    return net_G
+
+net_G = build_res_unet(n_input=1, n_output=2, size=256)
+
+class GANLoss(nn.Module):
+    def __init__(self, gan_mode='vanilla', real_label=1.0, fake_label=0.0):
+        super().__init__()
+        self.register_buffer('real_label', torch.tensor(real_label))
+        self.register_buffer('fake_label', torch.tensor(fake_label))
+        if gan_mode == 'vanilla':
+            self.loss = nn.BCEWithLogitsLoss()
+        elif gan_mode == 'lsgan':
+            self.loss = nn.MSELoss()
+
+    def get_labels(self, preds, target_is_real):
+        if target_is_real:
+            labels = self.real_label
+        else:
+            labels = self.fake_label
+        return labels.expand_as(preds)
+
+    def __call__(self, preds, target_is_real):
+        labels = self.get_labels(preds, target_is_real)
+        loss = self.loss(preds, labels)
+        return loss
+
+def load_model(model_class, file_path):
+    model = model_class(net_G=net_G)
+    model.load_state_dict(torch.load(file_path, map_location=device))
+
+    resnet_weights = torch.load(file_path)
+    resnet_weights = torch.load(r"model\res18-unet.pt")
+    resnet_state_dict = resnet_weights['state_dict'] if 'state_dict' in resnet_weights else resnet_weights
+
+    model_dict = model.state_dict()
+    filtered_resnet_state_dict = {k: v for k, v in resnet_state_dict.items() if k in model_dict}
+    model_dict.update(filtered_resnet_state_dict)
+    model.load_state_dict(model_dict)
+    return model
+
+    # return model
+    # model = model_class()
+    # model.load_state_dict(torch.load(file_path))
+    # return model
+
+def predict_color(model, image):
+    # img = Image.open(image)
+    img = image.resize((256, 256))
+    # to make it between -1 and 1
+    img = transforms.ToTensor()(img)[:1] * 2. - 1.
+
+    genimg = predict_and_return_image(model, img)
+    return genimg
+
+def predict_and_return_image(model, img):
+    model.eval()
+    with torch.no_grad():
+        preds = model.net_G(img.unsqueeze(0).to(device))
+    colorized = lab_to_rgb(img.unsqueeze(0), preds.cpu())[0]
+    return colorized

model.py

@@ -0,0 +1,142 @@
+import torch
+from torch import nn, optim
+from torchvision import transforms
+from torchvision.utils import make_grid
+from inference import init_model, GANLoss
+
+
+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)
+
+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)
+
+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/model_final.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:58876849eeea903233b5d0931ed0accabc5dd4230e5b897aa9aa0097df5ab93a
+size 135588892