added other files

app.py

@@ -0,0 +1,58 @@
+import gradio as gr
+import torch
+import matplotlib
+matplotlib.use('Agg')
+import numpy as np
+from PIL import Image
+import albumentations as A
+import albumentations.pytorch as al_pytorch
+import matplotlib.pyplot as plt
+import torchvision
+
+from app.model.lit_model import Pix2PixLitModule
+
+""" Load the model """
+model_checkpoint_path = "model/pix2pix_lightning_model/version_0/checkpoints/epoch=9-step=17780.ckpt"
+model = Pix2PixLitModule.load_from_checkpoint(
+    model_checkpoint_path
+)
+model.eval()
+
+
+def greet(name):
+    return "Hello " + name + "!!"
+
+
+def predict(image: Image):
+    # 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=np.asarray(image)
+    )['image'].unsqueeze(0)
+    result = model(inference_img)
+    result_grid = torchvision.utils.make_grid(
+            [result[0].permute(1, 2, 0).detach()],
+            normalize=True
+        )
+    plt.imsave("coloured_grid.png", result_grid.numpy())
+    torchvision.utils.save_image(result, "coloured_image.png", normalize=True)
+    return 'coloured_image.png', 'coloured_grid.png'
+
+
+if __name__ == '__main__':
+    #
+    iface = gr.Interface(
+        fn=predict,
+        inputs=gr.inputs.Image(type="pil"),
+        examples=["examples/thesis_test.png", "examples/thesis_test2.png"],
+        outputs=["image","image"],
+        title="Colour your sketches!",
+        description=" Upload a sketch and the conditional gan will colour it for you!",
+        article="WIP repo lives here - https://github.com/nmud19/thesisGAN "
+    )
+    iface.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

requirements.txt

@@ -0,0 +1,7 @@
+gradio
+torch
+torchvision
+pytorch_lightning
+matplotlib
+albumentations
+pillow