Working lightning bolts

app.py

@@ -1,6 +1,12 @@
+from typing import Union, List
+
 import gradio as gr
 import torch
+import torch.nn as nn
 import matplotlib
+import torch.nn.functional as F
+from pytorch_lightning.utilities.types import EPOCH_OUTPUT
+
 matplotlib.use('Agg')
 import numpy as np
 from PIL import Image
@@ -8,14 +14,153 @@ 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 """
+
+
+class OverpoweredPix2Pix(Pix2Pix):
+
+    def validation_step(self, batch, batch_idx):
+        """ Validation step """
+        real, condition = batch
+        with torch.no_grad():
+            loss = self._disc_step(real, condition)
+            self.log("val_PatchGAN_loss", loss)
+
+            loss = self._gen_step(real, condition)
+            self.log("val_generator_loss", loss)
+
+        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']
+        with torch.no_grad():
+            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)
+
+
+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.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/pix2pix_lightning_model/version_0/checkpoints/epoch=9-step=17780.ckpt"
-model = Pix2PixLitModule.load_from_checkpoint(
+model_checkpoint_path = "model/lightning_bolts_model/epoch=8-step=8010.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
 )
+
+model_chk = torch.load(
+    model_checkpoint_path, map_location=torch.device('cpu')
+)
+# model = gen().load_state_dict(model_chk)
+
 model.eval()
 
 
@@ -23,32 +168,57 @@ def greet(name):
     return "Hello " + name + "!!"
 
 
-def predict(image: Image):
+def predict(img: Image):
+    # 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=np.asarray(image)
+        image=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'
+    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"),
-    examples=["examples/thesis_test.png", "examples/thesis_test2.png"],
-    outputs=["image","image"],
+    #inputs="sketchpad",
+    examples=[
+        "examples/thesis_test.png",
+        "examples/thesis_test2.png",
+        # "examples/1000000.png"
+    ],
+    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 "