update gradio

app.py

@@ -5,7 +5,7 @@ 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
@@ -13,13 +13,17 @@ import albumentations.pytorch as al_pytorch
 import torchvision
 from pl_bolts.models.gans import Pix2Pix
 
+# Hack for spaces
+import os
+os.system("pip uninstall -y gradio")
+os.system("pip install -r requirements.txt")
+
 """ 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)
@@ -28,33 +32,36 @@ 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
+            f"Image Grid {str(self.current_epoch)}", grid_image, self.current_epoch
         )
 
 
 """ 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"
+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"
 
@@ -72,18 +79,20 @@ train_16_val_1_plbolts_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=[.5, .5, .5], std=[.5, .5, .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=[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)
 
     # Choose model
     if type_of_model == "train batch size 16, val batch size 1":
@@ -113,7 +122,7 @@ model_input = gr.inputs.Radio(
         "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 = [
@@ -132,13 +141,17 @@ with gr.Blocks() as demo:
     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(
+        " 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,
@@ -149,7 +162,9 @@ 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(

app/config.py

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

app/consume_data/consume_data.py

@@ -1,165 +0,0 @@
-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

@@ -1,69 +0,0 @@
-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

@@ -1,137 +0,0 @@
-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

@@ -1,174 +0,0 @@
-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

@@ -1,106 +0,0 @@
-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

@@ -1,145 +0,0 @@
-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

@@ -1,34 +0,0 @@
-
-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