Recreate the demo

app.py

@@ -1,10 +1,8 @@
 from typing import Union, List
 
 import gradio as gr
-import torch
-import torch.nn as nn
 import matplotlib
-import torch.nn.functional as F
+import torch
 from pytorch_lightning.utilities.types import EPOCH_OUTPUT
 
 matplotlib.use('Agg')
@@ -12,11 +10,8 @@ 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 pl_bolts.models.gans import Pix2Pix
-from app.generator.unetGen import Generator as gen
-from app.model.lit_model import Pix2PixLitModule
 
 """ Class """
 
@@ -49,182 +44,131 @@ class OverpoweredPix2Pix(Pix2Pix):
             ],
             normalize=True
         )
-        self.logger.experiment.add_image(f'Image Grid {str(self.current_epoch)}', grid_image, self.current_epoch)
-
-
-class Downsample(nn.Module):
-    def __init__(self, in_channels, out_channels):
-        super(Downsample, self).__init__()
-        self.conv_relu = nn.Sequential(
-            nn.Conv2d(in_channels, out_channels, 3, 2, 1),
-            nn.LeakyReLU(inplace=True)
-        )
-        self.bn = nn.BatchNorm2d(out_channels)
-
-    def forward(self, x, is_bn=True):
-        x = self.conv_relu(x)
-        if is_bn:
-            x = self.bn(x)
-        return x
-
-
-class Upsample(nn.Module):
-    def __init__(self, in_channels, out_channels):
-        super(Upsample, self).__init__()
-        self.upconv_relu = nn.Sequential(
-            nn.ConvTranspose2d(in_channels, out_channels, 3, 2, 1,
-                               output_padding=1),
-            nn.LeakyReLU(inplace=True)
+        self.logger.experiment.add_image(
+            f'Image Grid {str(self.current_epoch)}',
+            grid_image,
+            self.current_epoch
         )
-        self.bn = nn.BatchNorm2d(out_channels)
-
-    def forward(self, x, is_drop=False):
-        x = self.upconv_relu(x)
-        x = self.bn(x)
-        if is_drop:
-            x = F.dropout2d(x)
-        return x
-
-
-class Generator(nn.Module):
-    def __init__(self):
-        super(Generator, self).__init__()
-        self.down1 = Downsample(3, 64)
-        self.down2 = Downsample(64, 128)
-        self.down3 = Downsample(128, 256)
-        self.down4 = Downsample(256, 512)
-        self.down5 = Downsample(512, 512)
-        self.down6 = Downsample(512, 512)
-        self.down7 = Downsample(512, 512)
-        self.down8 = Downsample(512, 512)
-
-        self.up1 = Upsample(512, 512)
-        self.up2 = Upsample(1024, 512)
-        self.up3 = Upsample(1024, 512)
-        self.up4 = Upsample(1024, 512)
-        self.up5 = Upsample(1024, 256)
-        self.up6 = Upsample(512, 128)
-        self.up7 = Upsample(256, 64)
-
-        self.last = nn.ConvTranspose2d(128, 3,
-                                       kernel_size=3,
-                                       stride=2,
-                                       padding=1,
-                                       output_padding=1)
-
-    def forward(self, x):
-        x1 = self.down1(x)  # torch.Size([8, 64, 128, 128])
-        x2 = self.down2(x1)  # torch.Size([8, 128, 64, 64])
-        x3 = self.down3(x2)  # torch.Size([8, 256, 32, 32])
-        x4 = self.down4(x3)  # torch.Size([8, 512, 16, 16])
-        x5 = self.down5(x4)  # torch.Size([8, 512, 8, 8])
-        x6 = self.down6(x5)  # torch.Size([8, 512, 4, 4])
-        x7 = self.down7(x6)  # torch.Size([8, 512, 2, 2])
-        x8 = self.down8(x7)  # torch.Size([8, 512, 1, 1])
-
-        x8 = self.up1(x8, is_drop=True)  # torch.Size([8, 512, 2, 2])
-        x8 = torch.cat([x7, x8], dim=1)  # torch.Size([8, 1024, 2, 2])
-
-        x8 = self.up2(x8, is_drop=True)  # torch.Size([8, 512, 4, 4])
-        x8 = torch.cat([x6, x8], dim=1)  # torch.Size([8, 1024, 2, 2])
-
-        x8 = self.up3(x8, is_drop=True)  # torch.Size([8, 512, 8, 8])
-        x8 = torch.cat([x5, x8], dim=1)  # torch.Size([8, 1024, 8, 8])
-
-        x8 = self.up4(x8)  # torch.Size([8, 512, 16, 16])
-        x8 = torch.cat([x4, x8], dim=1)  # torch.Size([8, 1024, 16, 16])
-
-        x8 = self.up5(x8)
-        x8 = torch.cat([x3, x8], dim=1)
-
-        x8 = self.up6(x8)
-        x8 = torch.cat([x2, x8], dim=1)
-
-        x8 = self.up7(x8)
-        x8 = torch.cat([x1, x8], dim=1)
-
-        x8 = torch.tanh(self.last(x8))
-        return x8
 
 
 """ Load the model """
-model_checkpoint_path = "model/lightning_bolts_model/epoch=99-step=89000.ckpt"
+# 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"
 # model_checkpoint_path = "model/pix2pix_lightning_model/version_0/checkpoints/epoch=199-step=355600.ckpt"
 # model_checkpoint_path = "model/pix2pix_lightning_model/gen.pth"
 
-model = OverpoweredPix2Pix.load_from_checkpoint(
-    model_checkpoint_path
+# Load the models
+train_64_val_16_plbolts_model = OverpoweredPix2Pix.load_from_checkpoint(
+    train_64_val_16_plbolts_model_chkpt
 )
+train_64_val_16_plbolts_model.eval()
 
-model_chk = torch.load(
-    model_checkpoint_path, map_location=torch.device('cpu')
+#
+train_16_val_1_plbolts_model = OverpoweredPix2Pix.load_from_checkpoint(
+    train_16_val_1_plbolts_model_chkpt
 )
-# model = gen().load_state_dict(model_chk)
+train_16_val_1_plbolts_model.eval()
 
-model.eval()
 
-
-def greet(name):
-    return "Hello " + name + "!!"
-
-
-def predict(img: Image):
+def predict(img: Image, type_of_model: str):
+    """ Create predictions """
     # transform img
     image = np.asarray(img)
-    # image = image[:, image.shape[1] // 2:, :]
     # 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(),
     ])
-    # inverse_transform = A.Compose([
-    #     A.Normalize(
-    #         mean=[0.485, 0.456, 0.406],
-    #         std=[0.229, 0.224, 0.225]
-    #     ),
-    # ])
     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
+    else:
+        raise Exception("NOT YET SUPPORTED")
+
     with torch.no_grad():
         result = model.gen(inference_img)
-        # torchvision.utils.save_image(inference_img, "inference_image.png", normalize=True)
         torchvision.utils.save_image(result, "inference_image.png", normalize=True)
-
-    """
-    result_grid = torchvision.utils.make_grid(
-        [result[0]],
-        normalize=True
-    )
-    # plt.imsave("coloured_grid.png", (result_grid.permute(1,2,0).detach().numpy()*255).astype(int))
-    torchvision.utils.save_image(
-        result_grid, "coloured_image.png", normalize=True
-    )
-    """
     return "inference_image.png"  # 'coloured_image.png',
 
 
-iface = gr.Interface(
-    fn=predict,
-    inputs=gr.inputs.Image(type="pil"),
-    #inputs="sketchpad",
-    examples=[
-        "examples/thesis_test.png",
-        "examples/thesis_test2.png",
-        "examples/thesis1.png",
-        "examples/thesis4.png",
-        "examples/thesis5.png",
-        "examples/thesis6.png",
-        # "examples/1000000.png"
+def predict1(img: Image):
+    return predict(img=img, type_of_model="train batch size 16, val batch size 1")
+
+
+def predict2(img: Image):
+    return predict(img=img, type_of_model="train batch size 64, val batch size 16")
+
+
+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",
     ],
-    outputs=gr.outputs.Image(type="pil",),
-    #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 "
+    label="Type of Pix2Pix model to use : "
 )
-iface.launch()
+image_input = gr.inputs.Image(type="pil")
+img_examples = [
+    "examples/thesis_test.png",
+    "examples/thesis_test2.png",
+    "examples/thesis1.png",
+    "examples/thesis4.png",
+    "examples/thesis5.png",
+    "examples/thesis6.png",
+]
+
+
+with gr.Blocks() as demo:
+    gr.Markdown(" # Colour your sketches!")
+    gr.Markdown(" ## Description :")
+    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")
+    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", )
+            colour_1 = gr.Button("Colour it!")
+            gr.Examples(
+                examples=img_examples,
+                inputs=image_input1,
+                outputs=image_output1,
+                fn=predict1,
+            )
+        with gr.TabItem("tr_64_val_14"):
+            with gr.Row():
+                image_input2 = gr.inputs.Image(type="pil")
+                image_output2 = gr.outputs.Image(type="pil", )
+            colour_2 = gr.Button("Colour it!")
+            with gr.Row():
+                gr.Examples(
+                    examples=img_examples,
+                    inputs=image_input2,
+                    outputs=image_output2,
+                    fn=predict2,
+                )
+
+    colour_1.click(
+        fn=predict1,
+        inputs=image_input1,
+        outputs=image_output1,
+    )
+    colour_2.click(
+        fn=predict2,
+        inputs=image_input2,
+        outputs=image_output2,
+    )
+
+demo.title = "Colour your sketches!"
+demo.launch()