add new demo interface

app.py

@@ -1,11 +1,3 @@
-# Hack for spaces
-import os
-
-os.system("pip uninstall -y gradio")
-os.system("pip install -r requirements.txt")
-
-# Real code begins
-
 from typing import Union, List
 
 import gradio as gr
@@ -13,22 +5,21 @@ import matplotlib
 import torch
 from pytorch_lightning.utilities.types import EPOCH_OUTPUT
 
-matplotlib.use("Agg")
+matplotlib.use('Agg')
 import numpy as np
 from PIL import Image
 import albumentations as A
 import albumentations.pytorch as al_pytorch
 import torchvision
 from pl_bolts.models.gans import Pix2Pix
-from pl_bolts.models.gans.pix2pix.components import PatchGAN
-import torchvision.models as models
 
 """ Class """
 
 
 class OverpoweredPix2Pix(Pix2Pix):
+
     def validation_step(self, batch, batch_idx):
-        """Validation step"""
+        """ Validation step """
         real, condition = batch
         with torch.no_grad():
             loss = self._disc_step(real, condition)
@@ -37,56 +28,33 @@ class OverpoweredPix2Pix(Pix2Pix):
             loss = self._gen_step(real, condition)
             self.log("val_generator_loss", loss)
 
-        return {"sketch": real, "colour": condition}
+        return {
+            'sketch': real,
+            'colour': condition
+        }
 
-    def validation_epoch_end(
-        self, outputs: Union[EPOCH_OUTPUT, List[EPOCH_OUTPUT]]
-    ) -> None:
-        sketch = outputs[0]["sketch"]
-        colour = outputs[0]["colour"]
+    def validation_epoch_end(self, outputs: Union[EPOCH_OUTPUT, List[EPOCH_OUTPUT]]) -> None:
+        sketch = outputs[0]['sketch']
+        colour = outputs[0]['colour']
         with torch.no_grad():
             gen_coloured = self.gen(sketch)
         grid_image = torchvision.utils.make_grid(
             [
-                sketch[0],
-                colour[0],
-                gen_coloured[0],
+                sketch[0], colour[0], gen_coloured[0],
             ],
-            normalize=True,
+            normalize=True
         )
         self.logger.experiment.add_image(
-            f"Image Grid {str(self.current_epoch)}", grid_image, self.current_epoch
-        )
-
-
-class PatchGanChanged(OverpoweredPix2Pix):
-    def __init__(self, in_channels, out_channels):
-        super(PatchGanChanged, self).__init__(
-            in_channels=in_channels, out_channels=out_channels
-        )
-        self.patch_gan = self.get_dense_PatchGAN(self.patch_gan)
-
-    @staticmethod
-    def get_dense_PatchGAN(disc: PatchGAN) -> PatchGAN:
-        """Add final layer to gan"""
-        disc.final = torch.nn.Sequential(
-            disc.final,
-            torch.nn.Flatten(),
-            torch.nn.Linear(16 * 16, 1),
+            f'Image Grid {str(self.current_epoch)}',
+            grid_image,
+            self.current_epoch
         )
-        return disc
 
 
 """ Load the model """
 # train_64_val_16_patchgan_1val_plbolts_model_chkpt = "model/lightning_bolts_model/modified_path_gan.ckpt"
-train_64_val_16_plbolts_model_chkpt = (
-    "model/lightning_bolts_model/epoch=99-step=44600.ckpt"
-)
-train_16_val_1_plbolts_model_chkpt = (
-    "model/lightning_bolts_model/epoch=99-step=89000.ckpt"
-)
-modified_patchgan_chkpt = "model/lightning_bolts_model/modified_patchgan.ckpt"
-
+train_64_val_16_plbolts_model_chkpt = "model/lightning_bolts_model/epoch=99-step=44600.ckpt"
+train_16_val_1_plbolts_model_chkpt = "model/lightning_bolts_model/epoch=99-step=89000.ckpt"
 # model_checkpoint_path = "model/pix2pix_lightning_model/version_0/checkpoints/epoch=199-step=355600.ckpt"
 # model_checkpoint_path = "model/pix2pix_lightning_model/gen.pth"
 
@@ -102,142 +70,28 @@ train_16_val_1_plbolts_model = OverpoweredPix2Pix.load_from_checkpoint(
 )
 train_16_val_1_plbolts_model.eval()
 
-#
-modified_patchgan_model = PatchGanChanged.load_from_checkpoint(modified_patchgan_chkpt)
-modified_patchgan_model.eval()
-
-
-# Create new class
-class OverpoweredPix2Pix(Pix2Pix):
-    def __init__(self, in_channels, out_channels):
-        super(OverpoweredPix2Pix, self).__init__(
-            in_channels=in_channels, out_channels=out_channels
-        )
-        self._create_inception_score()
-
-    def _gen_step(self, real_images, conditioned_images):
-        # Pix2Pix has adversarial and a reconstruction loss
-        # First calculate the adversarial loss
-        fake_images = self.gen(conditioned_images)
-        disc_logits = self.patch_gan(fake_images, conditioned_images)
-        adversarial_loss = self.adversarial_criterion(
-            disc_logits, torch.ones_like(disc_logits)
-        )
-
-        # calculate reconstruction loss
-        recon_loss = self.recon_criterion(fake_images, real_images)
-        lambda_recon = self.hparams.lambda_recon
-
-        # calculate cosine similarity
-        representations_real = self.feature_extractor(real_images).flatten(1)
-        representations_fake = self.feature_extractor(fake_images).flatten(1)
-        similarity_score_list = self.cosine_similarity(
-            representations_real, representations_fake
-        )
-        cosine_sim = sum(similarity_score_list) / len(similarity_score_list)
-
-        self.log("Gen Cosine Sim Loss ", 1 - cosine_sim.cpu().detach().numpy())
-        # print(adversarial_loss,1-cosine_sim, lambda_recon, recon_loss, )
-
-        return (
-            (adversarial_loss)
-            + (lambda_recon * recon_loss)
-            + (lambda_recon * (1 - cosine_sim))
-        )
-
-    def _create_inception_score(self):
-        # init a pretrained resnet
-        backbone = models.resnet50(pretrained=True)
-        num_filters = backbone.fc.in_features
-        layers = list(backbone.children())[:-1]
-        self.feature_extractor = torch.nn.Sequential(*layers)
-        self.cosine_similarity = torch.nn.CosineSimilarity(dim=1, eps=1e-6)
-
-    def validation_step(self, batch, batch_idx):
-        """Validation step"""
-        real, condition = batch
-        with torch.no_grad():
-            disc_loss = self._disc_step(real, condition)
-            self.log("Valid PatchGAN Loss", disc_loss)
-
-            gan_loss = self._gen_step(real, condition)
-            self.log("Valid Generator Loss", gan_loss)
-
-            #
-            fake_images = self.gen(condition)
-            representations_real = self.feature_extractor(real).flatten(1)
-            representations_fake = self.feature_extractor(fake_images).flatten(1)
-            similarity_score_list = self.cosine_similarity(
-                representations_real, representations_fake
-            )
-            cosine_sim = sum(similarity_score_list) / len(similarity_score_list)
-
-            self.log("Valid Cosine Sim", cosine_sim)
-
-        return {"sketch": condition, "colour": real}
-
-    def validation_epoch_end(
-        self, outputs: Union[EPOCH_OUTPUT, List[EPOCH_OUTPUT]]
-    ) -> None:
-        sketch = outputs[0]["sketch"]
-        colour = outputs[0]["colour"]
-        self.feature_extractor.eval()
-        with torch.no_grad():
-            gen_coloured = self.gen(sketch)
-            representations_gen = self.feature_extractor(gen_coloured).flatten(1)
-            representations_fake = self.feature_extractor(colour).flatten(1)
-
-        similarity_score_list = self.cosine_similarity(
-            representations_gen, representations_fake
-        )
-        similarity_score = sum(similarity_score_list) / len(similarity_score_list)
-
-        grid_image = torchvision.utils.make_grid(
-            [
-                sketch[0],
-                colour[0],
-                gen_coloured[0],
-            ],
-            normalize=True,
-        )
-        self.logger.experiment.add_image(
-            f"Image Grid {str(self.current_epoch)} __ {str(similarity_score)} ",
-            grid_image,
-            self.current_epoch,
-        )
-
-
-cosine_sim_model_chkpt = "model/lightning_bolts_model/cosine_sim_model.ckpt"
-
-cosine_sim_model = OverpoweredPix2Pix.load_from_checkpoint(cosine_sim_model_chkpt)
-cosine_sim_model.eval()
-
 
 def predict(img: Image, type_of_model: str):
-    """Create predictions"""
+    """ Create predictions """
     # transform img
     image = np.asarray(img)
     # use on inference
-    inference_transform = A.Compose(
-        [
-            A.Resize(width=256, height=256),
-            A.Normalize(
-                mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_pixel_value=255.0
-            ),
-            al_pytorch.ToTensorV2(),
-        ]
-    )
-    inference_img = inference_transform(image=image)["image"].unsqueeze(0)
+    inference_transform = A.Compose([
+        A.Resize(width=256, height=256),
+        A.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5], max_pixel_value=255.0),
+        al_pytorch.ToTensorV2(),
+    ])
+    inference_img = inference_transform(
+        image=image
+    )['image'].unsqueeze(0)
 
     # Choose model
     if type_of_model == "train batch size 16, val batch size 1":
         model = train_16_val_1_plbolts_model
     elif type_of_model == "train batch size 64, val batch size 16":
         model = train_64_val_16_plbolts_model
-    elif type_of_model == "cosine similarity":
-        model = cosine_sim_model
     else:
-        model = modified_patchgan_model
+        raise Exception("NOT YET SUPPORTED")
 
     with torch.no_grad():
         result = model.gen(inference_img)
@@ -253,27 +107,13 @@ def predict2(img: Image):
     return predict(img=img, type_of_model="train batch size 64, val batch size 16")
 
 
-def predict3(img: Image):
-    return predict(
-        img=img,
-        type_of_model="train batch size 64, val batch size 16, patch gan has 1 output score instead of 16*16",
-    )
-
-
-def predict4(img: Image):
-    return predict(
-        img=img,
-        type_of_model="cosine similarity",
-    )
-
-
 model_input = gr.inputs.Radio(
     [
         "train batch size 16, val batch size 1",
         "train batch size 64, val batch size 16",
         "train batch size 64, val batch size 16, patch gan has 1 output score instead of 16*16",
     ],
-    label="Type of Pix2Pix model to use : ",
+    label="Type of Pix2Pix model to use : "
 )
 image_input = gr.inputs.Image(type="pil")
 img_examples = [
@@ -285,26 +125,20 @@ img_examples = [
     "examples/thesis6.png",
 ]
 
+
 with gr.Blocks() as demo:
     gr.Markdown(" # Colour your sketches!")
     gr.Markdown(" ## Description :")
-    gr.Markdown(" There are 4 Pix2Pix models in this example:")
+    gr.Markdown(" There are three Pix2Pix models in this example:")
     gr.Markdown(" 1. Training batch size is 16 , validation is 1")
     gr.Markdown(" 2. Training batch size is 64 , validation is 16")
-    gr.Markdown(
-        " 3. PatchGAN is changed, 1 value only instead of 16*16 ;"
-        "training batch size is 64 , validation is 16"
-    )
-    gr.Markdown(
-        " 4. cosine similarity is also added as a metric in this experiment for the generator. "
-    )
+    gr.Markdown(" 3. PatchGAN is changed, 1 value only instead of 16*16 ;"
+                "training batch size is 64 , validation is 16")
     with gr.Tabs():
         with gr.TabItem("tr_16_val_1"):
             with gr.Row():
                 image_input1 = gr.inputs.Image(type="pil")
-                image_output1 = gr.outputs.Image(
-                    type="pil",
-                )
+                image_output1 = gr.outputs.Image(type="pil", )
             colour_1 = gr.Button("Colour it!")
             gr.Examples(
                 examples=img_examples,
@@ -315,9 +149,7 @@ with gr.Blocks() as demo:
         with gr.TabItem("tr_64_val_14"):
             with gr.Row():
                 image_input2 = gr.inputs.Image(type="pil")
-                image_output2 = gr.outputs.Image(
-                    type="pil",
-                )
+                image_output2 = gr.outputs.Image(type="pil", )
             colour_2 = gr.Button("Colour it!")
             with gr.Row():
                 gr.Examples(
@@ -326,34 +158,6 @@ with gr.Blocks() as demo:
                     outputs=image_output2,
                     fn=predict2,
                 )
-        with gr.TabItem("Single Value Discriminator"):
-            with gr.Row():
-                image_input3 = gr.inputs.Image(type="pil")
-                image_output3 = gr.outputs.Image(
-                    type="pil",
-                )
-            colour_3 = gr.Button("Colour it!")
-            with gr.Row():
-                gr.Examples(
-                    examples=img_examples,
-                    inputs=image_input3,
-                    outputs=image_output3,
-                    fn=predict3,
-                )
-        with gr.TabItem("Cosine similarity loss"):
-            with gr.Row():
-                image_input4 = gr.inputs.Image(type="pil")
-                image_output4 = gr.outputs.Image(
-                    type="pil",
-                )
-            colour_4 = gr.Button("Colour it!")
-            with gr.Row():
-                gr.Examples(
-                    examples=img_examples,
-                    inputs=image_input4,
-                    outputs=image_output4,
-                    fn=predict4,
-                )
 
     colour_1.click(
         fn=predict1,
@@ -365,16 +169,6 @@ with gr.Blocks() as demo:
         inputs=image_input2,
         outputs=image_output2,
     )
-    colour_3.click(
-        fn=predict3,
-        inputs=image_input3,
-        outputs=image_output3,
-    )
-    colour_4.click(
-        fn=predict4,
-        inputs=image_input4,
-        outputs=image_output4,
-    )
 
 demo.title = "Colour your sketches!"
 demo.launch()

app/config.py

@@ -0,0 +1,3 @@
+num_workers = 4
+train_batch_size = 32
+val_batch_size = 1

app/consume_data/consume_data.py

@@ -0,0 +1,165 @@
+import torch
+import os
+from typing import List, Optional
+from PIL import Image
+import matplotlib.pyplot as plt
+from torchvision import transforms
+import albumentations as A
+import numpy as np
+import albumentations.pytorch as al_pytorch
+from typing import Dict, Tuple
+from app import config
+import pytorch_lightning as pl
+
+torch.__version__
+
+
+class AnimeDataset(torch.utils.data.Dataset):
+    """ Sketchs and Colored Image dataset """
+
+    def __init__(self, imgs_path: List[str], transforms: transforms.Compose) -> None:
+        """ Set the transforms and file path """
+        self.list_files = imgs_path
+        self.transform = transforms
+
+    def __len__(self) -> int:
+        """ Should return number of files """
+        return len(self.list_files)
+
+    def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        """ Get image and mask by index """
+        # read image file
+        img_file = self.list_files[index]
+        # img_path = os.path.join(self.root_dir, img_file)
+        image = np.array(Image.open(img_file))
+
+        # divide image into sketchs and colored_imgs, right is sketch and left is colored images
+        sketchs = image[:, image.shape[1] // 2:, :]
+        colored_imgs = image[:, :image.shape[1] // 2, :]
+
+        # data augmentation on both sketchs and colored_imgs
+        augmentations = self.transform.both_transform(image=sketchs, image0=colored_imgs)
+        sketchs, colored_imgs = augmentations['image'], augmentations['image0']
+
+        # conduct data augmentation respectively
+        sketchs = self.transform.transform_only_input(image=sketchs)['image']
+        colored_imgs = self.transform.transform_only_mask(image=colored_imgs)['image']
+        return sketchs, colored_imgs
+
+
+# Data Augmentation
+class Transforms:
+    def __init__(self):
+        # use on both sketchs and colored images
+        self.both_transform = A.Compose([
+            A.Resize(width=256, height=256),
+            A.HorizontalFlip(p=.5)
+        ], additional_targets={'image0': 'image'})
+
+        # use on sketchs only
+        self.transform_only_input = A.Compose([
+            A.ColorJitter(p=.1),
+            A.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5], max_pixel_value=255.0),
+            al_pytorch.ToTensorV2(),
+        ])
+
+        # use on colored images
+        self.transform_only_mask = A.Compose([
+            A.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5], max_pixel_value=255.0),
+            al_pytorch.ToTensorV2(),
+        ])
+
+
+class Transforms_v1:
+    """ Class to hold transforms """
+
+    def __init__(self):
+        # use on both sketchs and colored images
+        self.resize_572 = A.Compose([
+            A.Resize(width=572, height=572)
+        ])
+
+        self.resize_388 = A.Compose([
+            A.Resize(width=388, height=388)
+        ])
+
+        self.resize_256 = A.Compose([
+            A.Resize(width=256, height=256)
+        ])
+
+        # use on sketchs only
+        self.transform_only_input = A.Compose([
+            # A.ColorJitter(p=.1),
+            A.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5], max_pixel_value=255.0),
+            al_pytorch.ToTensorV2(),
+        ])
+
+        # use on colored images
+        self.transform_only_mask = A.Compose([
+            A.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5], max_pixel_value=255.0),
+            al_pytorch.ToTensorV2(),
+        ])
+
+
+class AnimeSketchDataModule(pl.LightningDataModule):
+    """ Class to hold the Anime sketch Data"""
+
+    def __init__(
+            self,
+            data_dir: str,
+            train_folder_name: str = "train/",
+            val_folder_name: str = "val/",
+            train_batch_size: int = config.train_batch_size,
+            val_batch_size: int = config.val_batch_size,
+            train_num_images: int = 0,
+            val_num_images: int = 0,
+    ):
+        super().__init__()
+        self.val_dataset = None
+        self.train_dataset = None
+        self.data_dir: str = data_dir
+        # Set train and val images folder
+        train_path: str = f"{self.data_dir}{train_folder_name}/"
+        train_images: List[str] = [f"{train_path}{x}" for x in os.listdir(train_path)]
+        val_path: str = f"{self.data_dir}{val_folder_name}"
+        val_images: List[str] = [f"{val_path}{x}" for x in os.listdir(val_path)]
+        #
+        self.train_images = train_images[:train_num_images] if train_num_images else train_images
+        self.val_images = val_images[:val_num_images] if val_num_images else val_images
+        #
+        self.train_batch_size = train_batch_size
+        self.val_batch_size = val_batch_size
+
+    def set_datasets(self) -> None:
+        """ Get the train and test datasets """
+        self.train_dataset = AnimeDataset(
+            imgs_path=self.train_images,
+            transforms=Transforms()
+        )
+        self.val_dataset = AnimeDataset(
+            imgs_path=self.val_images,
+            transforms=Transforms()
+        )
+        print("The train test dataset lengths are : ", len(self.train_dataset), len(self.val_dataset))
+        return None
+
+    def setup(self, stage: Optional[str] = None) -> None:
+        self.set_datasets()
+
+    def train_dataloader(self):
+        return torch.utils.data.DataLoader(
+            self.train_dataset,
+            batch_size=self.train_batch_size,
+            shuffle=False,
+            num_workers=2,
+            pin_memory=True
+        )
+
+    def val_dataloader(self):
+        return torch.utils.data.DataLoader(
+            self.val_dataset,
+            batch_size=self.val_batch_size,
+            shuffle=False,
+            num_workers=2,
+            pin_memory=True
+        )

app/data.py

@@ -0,0 +1,69 @@
+import torch
+import os
+from typing import List
+from PIL import Image
+import matplotlib.pyplot as plt
+from torchvision import transforms
+import albumentations as A
+import numpy as np
+import albumentations.pytorch as al_pytorch
+from typing import Dict, Tuple
+
+
+class AnimeDataset(torch.utils.data.Dataset):
+    """ Sketchs and Colored Image dataset """
+
+    def __init__(self, imgs_path: List[str], transforms: transforms.Compose) -> None:
+        """ Set the transforms and file path """
+        self.list_files = imgs_path
+        self.transform = transforms
+
+    def __len__(self) -> int:
+        """ Should return number of files """
+        return len(self.list_files)
+
+    def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        """ Get image and mask by index """
+        # read image file
+        img_path = img_file = self.list_files[index]
+        image = np.array(Image.open(img_path))
+
+        # divide image into sketchs and colored_imgs, right is sketch and left is colored images
+        # as according to the dataset
+        sketchs = image[:, image.shape[1] // 2:, :]
+        colored_imgs = image[:, :image.shape[1] // 2, :]
+
+        # data augmentation on both sketchs and colored_imgs
+        augmentations = self.transform.both_transform(image=sketchs, image0=colored_imgs)
+        sketchs, colored_imgs = augmentations['image'], augmentations['image0']
+
+        # conduct data augmentation respectively
+        sketchs = self.transform.transform_only_input(image=sketchs)['image']
+        colored_imgs = self.transform.transform_only_mask(image=colored_imgs)['image']
+        return sketchs, colored_imgs
+
+
+class Transforms:
+    """ Class to hold transforms """
+
+    def __init__(self):
+        # use on both sketchs and colored images
+        self.both_transform = A.Compose([
+            A.Resize(width=1024, height=1024),
+            A.HorizontalFlip(p=.5)
+        ],
+            additional_targets={'image0': 'image'}
+        )
+
+        # use on sketchs only
+        self.transform_only_input = A.Compose([
+            # A.ColorJitter(p=.1),
+            A.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5], max_pixel_value=255.0),
+            al_pytorch.ToTensorV2(),
+        ])
+
+        # use on colored images
+        self.transform_only_mask = A.Compose([
+            A.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5], max_pixel_value=255.0),
+            al_pytorch.ToTensorV2(),
+        ])

app/discriminator/patch_gan.py

@@ -0,0 +1,137 @@
+import torch.nn as nn
+import torch
+import albumentations as A
+
+
+# CNN block will be used repeatly later
+class CNNBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride=2):
+        super().__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 4, stride, bias=False, padding_mode='reflect'),
+            nn.BatchNorm2d(out_channels),
+            nn.LeakyReLU(0.2)
+        )
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class PatchGan(torch.nn.Module):
+    """ Patch GAN Architecture """
+
+    @staticmethod
+    def create_contracting_block(in_channels: int, out_channels: int):
+        """
+        Create encoding layer
+        :param in_channels:
+        :param out_channels:
+        :return:
+        """
+        conv_layer = torch.nn.Sequential(
+            torch.nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=3,
+                padding=1,
+            ),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(
+                in_channels=out_channels,
+                out_channels=out_channels,
+                kernel_size=3,
+                padding=1,
+            ),
+            torch.nn.ReLU(),
+        )
+        max_pool = torch.nn.Sequential(
+            torch.nn.MaxPool2d(
+                stride=2,
+                kernel_size=2,
+            ),
+        )
+        layer = torch.nn.Sequential(
+            conv_layer,
+            max_pool,
+        )
+        return layer
+
+    def __init__(self, input_channels: int, hidden_channels: int) -> None:
+        super().__init__()
+        self.resize_channels = torch.nn.Conv2d(
+            in_channels=input_channels,
+            out_channels=hidden_channels,
+            kernel_size=1,
+        )
+
+        self.enc1 = self.create_contracting_block(
+            in_channels=hidden_channels,
+            out_channels=hidden_channels * 2
+        )
+
+        self.enc2 = self.create_contracting_block(
+            in_channels=hidden_channels * 2,
+            out_channels=hidden_channels * 4
+        )
+
+        self.enc3 = self.create_contracting_block(
+            in_channels=hidden_channels * 4,
+            out_channels=hidden_channels * 8
+        )
+        self.enc4 = self.create_contracting_block(
+            in_channels=hidden_channels * 8,
+            out_channels=hidden_channels * 16
+        )
+
+        self.final_layer = torch.nn.Conv2d(
+            in_channels=hidden_channels * 16,
+            out_channels=1,
+            kernel_size=1,
+        )
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        """ Forward patch gan layer """
+        inpt = torch.cat([x, y], axis=1)
+        resize_img = self.resize_channels(inpt)
+        enc1 = self.enc1(resize_img)
+        enc2 = self.enc2(enc1)
+        enc3 = self.enc3(enc2)
+        enc4 = self.enc4(enc3)
+        final_layer = self.final_layer(enc4)
+        return final_layer
+
+
+# x, y <- concatenate the gen image and the input image to determin the gen image is real or not
+class Discriminator(nn.Module):
+    def __init__(self, in_channels=3, features=[64, 128, 256, 512]):
+        super().__init__()
+        self.initial = nn.Sequential(
+            nn.Conv2d(in_channels * 2, features[0], kernel_size=4, stride=2, padding=1, padding_mode='reflect'),
+            nn.LeakyReLU(.2)
+        )
+
+        # save layers into a list
+        layers = []
+        in_channels = features[0]
+        for feature in features[1:]:
+            layers.append(
+                CNNBlock(
+                    in_channels,
+                    feature,
+                    stride=1 if feature == features[-1] else 2
+                ),
+            )
+            in_channels = feature
+
+        # append last conv layer
+        layers.append(
+            nn.Conv2d(in_channels, 1, kernel_size=4, stride=1, padding=1, padding_mode='reflect')
+        )
+
+        # create a model using the list of layers
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x, y):
+        x = torch.cat([x, y], dim=1)
+        x = self.initial(x)
+        return self.model(x)

app/generator/unetGen.py

@@ -0,0 +1,174 @@
+import torch
+import torch.nn as nn
+from app.generator import unetParts
+
+
+class UNET(torch.nn.Module):
+    """ Implementation of unet """
+
+    def __init__(
+            self,
+    ) -> None:
+        """
+        Create the UNET here
+        """
+        super().__init__()
+        self.enc_layer1: unetParts.EncoderLayer = unetParts.EncoderLayer(
+            in_channels=3,
+            out_channels=64
+        )
+        self.enc_layer2: unetParts.EncoderLayer = unetParts.EncoderLayer(
+            in_channels=64,
+            out_channels=128
+        )
+        self.enc_layer3: unetParts.EncoderLayer = unetParts.EncoderLayer(
+            in_channels=128,
+            out_channels=256
+        )
+        self.enc_layer4: unetParts.EncoderLayer = unetParts.EncoderLayer(
+            in_channels=256,
+            out_channels=512
+        )
+        # Middle layer
+        self.middle_layer: unetParts.MiddleLayer = unetParts.MiddleLayer(
+            in_channels=512,
+            out_channels=1024,
+        )
+        # Decoding layer
+        self.dec_layer1: unetParts.DecoderLayer = unetParts.DecoderLayer(
+            in_channels=1024,
+            out_channels=512,
+        )
+        self.dec_layer2: unetParts.DecoderLayer = unetParts.DecoderLayer(
+            in_channels=512,
+            out_channels=256,
+        )
+
+        self.dec_layer3: unetParts.DecoderLayer = unetParts.DecoderLayer(
+            in_channels=256,
+            out_channels=128,
+        )
+        self.dec_layer4: unetParts.DecoderLayer = unetParts.DecoderLayer(
+            in_channels=128,
+            out_channels=64,
+        )
+        self.final_layer: torch.nn.Conv2d = torch.nn.Conv2d(
+            in_channels=64,
+            out_channels=3,
+            kernel_size=1
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward function
+        :param x:
+        :return:
+        """
+        # enc layers
+        enc1, conv1 = self.enc_layer1(x=x)  # 64
+        enc2, conv2 = self.enc_layer2(x=enc1)  # 128
+        enc3, conv3 = self.enc_layer3(x=enc2)  # 256
+        enc4, conv4 = self.enc_layer4(x=enc3)  # 512
+        # middle layers
+        mid = self.middle_layer(x=enc4)  # 1024
+        # expanding layers
+        # 512
+        dec1 = self.dec_layer1(
+            input_layer=mid,
+            cropping_layer=conv4,
+        )
+        # 256
+        dec2 = self.dec_layer2(
+            input_layer=dec1,
+            cropping_layer=conv3,
+        )
+        # 128
+        dec3 = self.dec_layer3(
+            input_layer=dec2,
+            cropping_layer=conv2,
+        )
+        # 64
+        dec4 = self.dec_layer4(
+            input_layer=dec3,
+            cropping_layer=conv1,
+        )
+        # 3
+        fin_layer = self.final_layer(
+            dec4,
+        )
+        # Interpolate to retain size
+        fin_layer_resized = torch.nn.functional.interpolate(fin_layer, 572)
+        return fin_layer_resized
+
+
+class Generator(nn.Module):
+    def __init__(self, in_channels=3, features=64):
+        super().__init__()
+        # Encoder
+        self.initial_down = nn.Sequential(
+            nn.Conv2d(in_channels, features, 4, 2, 1, padding_mode='reflect'),
+            nn.LeakyReLU(.2),
+        )
+        self.down1 = Block(features, features * 2, down=True, act='leaky', use_dropout=False)  # 64
+        self.down2 = Block(features * 2, features * 4, down=True, act='leaky', use_dropout=False)  # 32
+        self.down3 = Block(features * 4, features * 8, down=True, act='leaky', use_dropout=False)  # 16
+        self.down4 = Block(features * 8, features * 8, down=True, act='leaky', use_dropout=False)  # 8
+        self.down5 = Block(features * 8, features * 8, down=True, act='leaky', use_dropout=False)  # 4
+        self.down6 = Block(features * 8, features * 8, down=True, act='leaky', use_dropout=False)  # 2
+        self.bottleneck = nn.Sequential(
+            nn.Conv2d(features * 8, features * 8, 4, 2, 1, padding_mode='reflect'),
+            nn.ReLU(),  # 1x1
+        )
+        # Decoder
+        self.up1 = Block(features * 8, features * 8, down=False, act='relu', use_dropout=True)
+        self.up2 = Block(features * 8 * 2, features * 8, down=False, act='relu', use_dropout=True)
+        self.up3 = Block(features * 8 * 2, features * 8, down=False, act='relu', use_dropout=True)
+        self.up4 = Block(features * 8 * 2, features * 8, down=False, act='relu', use_dropout=False)
+        self.up5 = Block(features * 8 * 2, features * 4, down=False, act='relu', use_dropout=False)
+        self.up6 = Block(features * 4 * 2, features * 2, down=False, act='relu', use_dropout=False)
+        self.up7 = Block(features * 2 * 2, features, down=False, act='relu', use_dropout=False)
+        self.final_up = nn.Sequential(
+            nn.ConvTranspose2d(features * 2, in_channels, kernel_size=4, stride=2, padding=1),
+            nn.Tanh()
+        )
+
+    def forward(self, x):
+        # Encoder
+        d1 = self.initial_down(x)
+        d2 = self.down1(d1)
+        d3 = self.down2(d2)
+        d4 = self.down3(d3)
+        d5 = self.down4(d4)
+        d6 = self.down5(d5)
+        d7 = self.down6(d6)
+        bottleneck = self.bottleneck(d7)
+
+        # Decoder
+        u1 = self.up1(bottleneck)
+        u2 = self.up2(torch.cat([u1, d7], 1))
+        u3 = self.up3(torch.cat([u2, d6], 1))
+        u4 = self.up4(torch.cat([u3, d5], 1))
+        u5 = self.up5(torch.cat([u4, d4], 1))
+        u6 = self.up6(torch.cat([u5, d3], 1))
+        u7 = self.up7(torch.cat([u6, d2], 1))
+        return self.final_up(torch.cat([u7, d1], 1))
+
+
+# block will be use repeatly later
+class Block(nn.Module):
+    def __init__(self, in_channels, out_channels, down=True, act='relu', use_dropout=False):
+        super().__init__()
+        self.conv = nn.Sequential(
+            # the block will be use on both encoder (down=True) and decoder (down=False)
+            nn.Conv2d(in_channels, out_channels, 4, 2, 1, bias=False, padding_mode='reflect')
+            if down
+            else nn.ConvTranspose2d(in_channels, out_channels, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU() if act == 'relu' else nn.LeakyReLU(.2)
+        )
+        self.use_dropout = use_dropout
+        self.dropout = nn.Dropout(.5)
+
+    def forward(self, x):
+        x = self.conv(x)
+        return self.dropout(x) if self.use_dropout else x

app/generator/unetParts.py

@@ -0,0 +1,106 @@
+import torch
+from typing import Tuple
+
+
+class DecoderLayer(torch.nn.Module):
+    """Decoder model"""
+
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        self.up_sample_layer = torch.nn.Sequential(
+            torch.nn.ConvTranspose2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=2,
+                stride=2,
+                bias=False,
+            )
+        )
+        self.conv_layer = EncoderLayer(
+            in_channels=in_channels,
+            out_channels=out_channels,
+        ).conv_layer
+
+    @staticmethod
+    def _get_cropping_shape(previous_layer_shape: torch.Size, current_layer_shape: torch.Size) -> int:
+        """ Get the shape to crop """
+        return (previous_layer_shape[2] - current_layer_shape[2]) // 2 * -1
+
+    def forward(
+            self,
+            input_layer: torch.Tensor,
+            cropping_layer: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Forward function to concatenate and conv the figure
+        :param cropping_layer:
+        :param input_layer:
+        :return:
+        """
+        input_layer = self.up_sample_layer(input_layer)
+
+        cropping_shape = self._get_cropping_shape(
+            current_layer_shape=input_layer.shape,
+            previous_layer_shape=cropping_layer.shape,
+        )
+
+        cropping_layer = torch.nn.functional.pad(
+            input=cropping_layer,
+            pad=[cropping_shape for _ in range(4)]
+        )
+        combined_layer = torch.cat(
+            tensors=[input_layer, cropping_layer],
+            dim=1
+        )
+        result = self.conv_layer(combined_layer)
+        return result
+
+
+class EncoderLayer(torch.nn.Module):
+    """Encoder Layer"""
+
+    def __init__(self, in_channels: int, out_channels: int) -> None:
+        super().__init__()
+        self.conv_layer = torch.nn.Sequential(
+            torch.nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=3,
+                stride=2,
+                padding=1,
+            ),
+            torch.nn.LeakyReLU(),
+            torch.nn.Conv2d(
+                in_channels=out_channels,
+                out_channels=out_channels,
+                kernel_size=3,
+                stride=2,
+                padding=1,
+            ),
+            torch.nn.LeakyReLU(),
+        )
+        self.max_pool = torch.nn.Sequential(
+            torch.nn.MaxPool2d(2),
+        )
+        self.layer = torch.nn.Sequential(
+            self.conv_layer,
+            self.max_pool,
+        )
+
+    def get_conv_layers(self, x: torch.Tensor) -> torch.Tensor:
+        """Need to concatenate the layer"""
+        return self.conv_layer(x)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass to return conv layer and the max pool layer"""
+        conv_output: torch.tensor = self.conv_layer(x)
+        fin_out: torch.Tensor = self.max_pool(conv_output)
+        return fin_out, conv_output
+
+
+class MiddleLayer(EncoderLayer):
+    """Middle layer only"""
+
+    def forward(self, x: torch.tensor) -> torch.tensor:
+        """Forward pass"""
+        return self.conv_layer(x)

app/model/lit_model.py

@@ -0,0 +1,145 @@
+import matplotlib.pyplot as plt
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+import torchvision
+
+
+class Pix2PixLitModule(pl.LightningModule):
+    """ Lightning Module for pix2pix """
+
+    @staticmethod
+    def _weights_init(m):
+        if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
+            torch.nn.init.normal_(m.weight, 0.0, 0.02)
+        if isinstance(m, nn.BatchNorm2d):
+            torch.nn.init.normal_(m.weight, 0.0, 0.02)
+            torch.nn.init.constant_(m.bias, 0)
+
+    def __init__(
+            self,
+            generator,
+            discriminator,
+            use_gpu: bool,
+            lambda_recon=100
+    ):
+        super().__init__()
+        self.save_hyperparameters()
+
+        self.gen = generator
+        self.disc = discriminator
+
+        # intializing weights
+        self.gen = self.gen.apply(self._weights_init)
+        self.disc = self.disc.apply(self._weights_init)
+        #
+        self.adversarial_criterion = nn.BCEWithLogitsLoss()
+        self.recon_criterion = nn.L1Loss()
+        self.lambda_l1 = lambda_recon
+
+    def _gen_step(self, sketch, coloured_sketches):
+        # Pix2Pix has adversarial and a reconstruction loss
+        # First calculate the adversarial loss
+        gen_coloured_sketches = self.gen(sketch)
+        # disc_logits = self.disc(gen_coloured_sketches, coloured_sketches)
+        disc_logits = self.disc(sketch, gen_coloured_sketches)
+        adversarial_loss = self.adversarial_criterion(disc_logits, torch.ones_like(disc_logits))
+        # calculate reconstruction loss
+        recon_loss = self.recon_criterion(gen_coloured_sketches, coloured_sketches) * self.lambda_l1
+        #
+        self.log("Gen recon_loss", recon_loss)
+        self.log("Gen adversarial_loss", adversarial_loss)
+        #
+        return adversarial_loss + recon_loss
+
+    def _disc_step(self, sketch, coloured_sketches):
+        gen_coloured_sketches = self.gen(sketch).detach()
+        #
+        # fake_logits = self.disc(gen_coloured_sketches, coloured_sketches)
+        fake_logits = self.disc(sketch, gen_coloured_sketches)
+        real_logits = self.disc(sketch, coloured_sketches)
+        #
+        fake_loss = self.adversarial_criterion(fake_logits, torch.zeros_like(fake_logits))
+        real_loss = self.adversarial_criterion(real_logits, torch.ones_like(real_logits))
+        #
+        self.log("PatchGAN fake_loss", fake_loss)
+        self.log("PatchGAN real_loss", real_loss)
+        return (real_loss + fake_loss) / 2
+
+    def forward(self, x):
+        return self.gen(x)
+
+    def training_step(self, batch, batch_idx, optimizer_idx):
+        real, condition = batch
+        loss = None
+        if optimizer_idx == 0:
+            loss = self._disc_step(real, condition)
+            self.log("TRAIN_PatchGAN Loss", loss)
+        elif optimizer_idx == 1:
+            loss = self._gen_step(real, condition)
+            self.log("TRAIN_Generator Loss", loss)
+        return loss
+
+    def validation_epoch_end(self, outputs) -> None:
+        """ Log the images"""
+        sketch = outputs[0]['sketch']
+        colour = outputs[0]['colour']
+        gen_coloured = self.gen(sketch)
+        grid_image = torchvision.utils.make_grid(
+            [sketch[0], colour[0], gen_coloured[0]],
+            normalize=True
+        )
+        self.logger.experiment.add_image(f'Image Grid {str(self.current_epoch)}', grid_image, self.current_epoch)
+        #plt.imshow(grid_image.permute(1, 2, 0))
+
+    def validation_step(self, batch, batch_idx):
+        """ Validation step """
+        real, condition = batch
+        return {
+            'sketch': real,
+            'colour': condition
+        }
+
+    def configure_optimizers(self, lr=2e-4):
+        gen_opt = torch.optim.Adam(self.gen.parameters(), lr=lr, betas=(0.5, 0.999))
+        disc_opt = torch.optim.Adam(self.disc.parameters(), lr=lr, betas=(0.5, 0.999))
+        return disc_opt, gen_opt
+
+# class EpochInference(pl.Callback):
+#     """
+#     Callback on each end of training epoch
+#     The callback will do inference on test dataloader based on corresponding checkpoints
+#     The results will be saved as an image with 4-rows:
+#         1 - Input image e.g. grayscale edged input
+#         2 - Ground-truth
+#         3 - Single inference
+#         4 - Mean of hundred accumulated inference
+#     Note that the inference have a noise factor that will generate different output on each execution
+#     """
+#
+#     def __init__(self, dataloader, use_gpu: bool, *args, **kwargs):
+#         super().__init__(*args, **kwargs)
+#         self.dataloader = dataloader
+#         self.use_gpu = use_gpu
+#
+#     def on_train_epoch_end(self, trainer, pl_module):
+#         super().on_train_epoch_end(trainer, pl_module)
+#         data = next(iter(self.dataloader))
+#         image, target = data
+#         if self.use_gpu:
+#             image = image.cuda()
+#             target = target.cuda()
+#         with torch.no_grad():
+#             # Take average of multiple inference as there is a random noise
+#             # Single
+#             reconstruction_init = pl_module(image)
+#             reconstruction_init = torch.clip(reconstruction_init, 0, 1)
+#             # # Mean
+#             # reconstruction_mean = torch.stack([pl_module(image) for _ in range(10)])
+#             # reconstruction_mean = torch.clip(reconstruction_mean, 0, 1)
+#             # reconstruction_mean = torch.mean(reconstruction_mean, dim=0)
+#         # Grayscale 1-D to 3-D
+#         # image = torch.stack([image for _ in range(3)], dim=1)
+#         # image = torch.squeeze(image)
+#         grid_image = torchvision.utils.make_grid([image[0], target[0], reconstruction_init[0]])
+#         torchvision.utils.save_image(grid_image, fp=f'{trainer.default_root_dir}/epoch-{trainer.current_epoch:04}.png')

app/scratch.py

@@ -0,0 +1,34 @@
+
+class GANInference:
+    def __init__(
+            self,
+            model: Pix2PixLitModule,
+            img_file: str = "/Users/nimud/Downloads/thesis_test2.png",
+    ) -> None:
+        self.img_file = img_file
+        self.model = model
+
+    def _get_image_from_path(self) -> torch.Tensor:
+        """ gets the tensor from filepath """
+        image = np.array(Image.open(self.img_file))
+        # use on inference
+        inference_transform = A.Compose([
+            A.Resize(width=256, height=256),
+            A.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5], max_pixel_value=255.0),
+            al_pytorch.ToTensorV2(),
+        ])
+        inference_img = inference_transform(image=image)['image'].unsqueeze(0)
+        return inference_img
+
+    def _create_grid(self, result: torch.Tensor) -> np.array:
+        return torchvision.utils.make_grid(
+            [result[0].permute(1, 2, 0).detach()],
+            normalize=True
+        )
+
+    def run(self) -> np.array:
+        """ Returns a plottable image """
+        inference_img = self._get_image_from_path()
+        result = self.model(inference_img)
+        adjusted_result = self._create_grid(result=result)
+        return adjusted_result

model/lightning_bolts_model/cosine_sim_model.ckpt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:2987394cad6890877faaf61ade50eada5397c2d1447a48049e8ad3197ea461cc
-size 780630439

model/lightning_bolts_model/modified_patchgan.ckpt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:7b6b85940399eb68eca7a62b603cd62ac2bc813bbec70a16df83842da73dd14a
-size 686280151