Update app.py

app.py

@@ -1,3 +1,130 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+from PIL import Image
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+import os
+from tqdm import tqdm
+from torchvision.utils import save_image
+import gradio as gr
+
+class cnnBlock(nn.Module):
+  def __init__(self, in_channels, out_channels, up_sample=False, use_act=True, **kwargs):
+    super().__init__()
+    self.cnn_block = nn.Sequential(
+        nn.ConvTranspose2d(in_channels=in_channels, out_channels=out_channels, **kwargs)
+        if up_sample else
+        nn.Conv2d(in_channels=in_channels, out_channels=out_channels, padding_mode="reflect", **kwargs),
+        nn.InstanceNorm2d(out_channels),
+        nn.ReLU(inplace=True) if use_act else nn.Identity()
+    )
+
+  def forward(self, x):
+    return self.cnn_block(x)
+
+class residualBlock(nn.Module):
+  def __init__(self, channels):
+    super().__init__()
+    self.resBlock = nn.Sequential(
+        cnnBlock(channels, channels, kernel_size=3, padding=1),
+        cnnBlock(channels, channels, use_act=False, kernel_size=3, padding=1)
+    )
+
+  def forward(self, x):
+    return x + self.resBlock(x)
+
+class Generator(nn.Module):
+  def __init__(self, img_channels=3, features=64, num_residual=9):
+    super().__init__()
+    self.initial = nn.Sequential(
+        nn.Conv2d(img_channels, 64, kernel_size=7, stride=1, padding=3, padding_mode="reflect"),
+        nn.ReLU()
+    )
+    self.downBlock = nn.ModuleList([
+        cnnBlock(features, features*2, kernel_size=3, stride=2, padding=1),
+        cnnBlock(features*2, features*4, kernel_size=3, stride=2, padding=1)
+    ])
+    self.resBlock = nn.Sequential(*[residualBlock(features*4) for _ in range(num_residual)])
+    self.upBlock = nn.ModuleList([
+        cnnBlock(features*4, features*2, up_sample=True, kernel_size=3, stride=2, padding=1, output_padding=1),
+        cnnBlock(features*2, features, up_sample=True, kernel_size=3, stride=2, padding=1, output_padding=1),
+    ])
+    self.final = nn.Conv2d(features, img_channels, kernel_size=7, stride=1, padding=3, padding_mode="reflect")
+
+  def forward(self, x):
+    x = self.initial(x)
+    for layer in self.downBlock:
+      x = layer(x)
+    x = self.resBlock(x)
+    for layer in self.upBlock:
+      x = layer(x)
+    x = self.final(x)
+    return torch.tanh(x)
+      
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+TRAIN_DIR = "/kaggle/input/vangogh2photo/vangogh2photo/train"
+VAL_DIR = "/kaggle/input/vangogh2photo/vangogh2photo/val"
+BATCH_SIZE = 1
+LEARNING_RATE = 2e-4
+LAMBDA_IDENTITY = 0.0
+LAMBDA_CYCLE = 10
+NUM_WORKERS = 4
+NUM_EPOCHS = 0
+LOAD_MODEL = True
+SAVE_MODEL = False
+CHECKPOINT_GEN_B = "/kaggle/input/checkpoints/genB.pth.tar"
+CHECKPOINT_GEN_A = "/kaggle/input/checkpoints/genA.pth.tar"
+CHECKPOINT_DISC_A = "/kaggle/input/checkpoints/discA.pth.tar"
+CHECKPOINT_DISC_B = "/kaggle/input/checkpoints/discB.pth.tar"
+
+transforms = A.Compose(
+    [
+        A.Resize(width=256, height=256),
+        A.HorizontalFlip(p=0.5),
+        A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_pixel_value=255),
+        ToTensorV2(),
+    ],
+    additional_targets={"image0": "image"},
+    is_check_shapes=False
+)
+
+def load_checkpoint(checkpoint_file, model, optimizer, lr):
+    print("=> Loading checkpoint")
+    checkpoint = torch.load(checkpoint_file, map_location=DEVICE)
+    model.load_state_dict(checkpoint["state_dict"])
+    optimizer.load_state_dict(checkpoint["optimizer"])
+
+    # If we don't do this then it will just have learning rate of old checkpoint
+    # and it will lead to many hours of debugging \:
+    for param_group in optimizer.param_groups:
+        param_group["lr"] = lr
+
+genA = Generator().to(DEVICE)
+
+load_checkpoint(CHECKPOINT_GEN_A, genA, optim_gen, LEARNING_RATE)
+
+def postprocess_and_show(output):
+    # Detach from GPU, move to CPU, and remove the batch dimension
+    output = output.squeeze(0).detach().cpu()  
+    
+    # Convert from [-1, 1] to [0, 1]
+    output = (output + 1) / 2.0  
+    
+    # Convert from tensor to NumPy array and transpose (C, H, W) to (H, W, C)
+    output_image = output.permute(1, 2, 0).numpy()
+    
+    # Convert to a [0, 255] image (optional if you're using a visualization library)
+    output_image = (output_image * 255).astype(np.uint8)
+    
+    # Option 2: Convert to a PIL image if you want to save or manipulate it
+    output_pil = Image.fromarray(output_image)
+
+    return output_pil
+    #plt.imshow(output_pil)
+
 transforms2 = A.Compose(
     [
         A.Resize(width=256, height=256),