Create app.py

app.py

@@ -0,0 +1,110 @@
+import gradio as gr
+import torch
+from PIL import Image
+from torchvision import transforms
+import torchvision.models as models
+import torch.nn as nn
+
+class DummyModel(nn.Module):
+    def __init__(self):
+        super(DummyModel, self).__init__()
+
+        self.encoder1 = nn.Sequential(
+            nn.Conv2d(1, 64, 3, 2, 1),  # 150x150 -> 75x75
+            nn.LeakyReLU()
+        )
+
+        self.encoder2 = nn.Sequential(
+            nn.Conv2d(64, 128, 3, 2, 1),  # 75x75 -> 38x38
+            nn.LeakyReLU()
+        )
+
+        self.encoder3 = nn.Sequential(
+            nn.Conv2d(128, 256, 3, 2, 1),  # 38x38 -> 19x19
+            nn.LeakyReLU()
+        )
+
+        self.encoder4 = nn.Sequential(
+            nn.Conv2d(256, 512, 3, 2, 1),  # 19x19 -> 10x10
+            nn.LeakyReLU()
+        )
+
+        # Bottleneck
+        self.bottleneck = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(512 * 10 * 10, 2048)
+        )
+
+        # Decoder
+        self.decoder_fc = nn.Sequential(
+            nn.Linear(2048, 512 * 10 * 10),
+            nn.Unflatten(1, (512, 10, 10))
+        )
+
+        self.decoder1 = nn.Sequential(
+            nn.ConvTranspose2d(512, 256, 3, 2, 1),  # 10x10 -> 19x19
+            nn.LeakyReLU()
+        )
+
+        self.decoder2 = nn.Sequential(
+            nn.ConvTranspose2d(256, 128, 3, 2, 1, output_padding=1),  # 19x19 -> 38x38
+            nn.LeakyReLU()
+        )
+
+        self.decoder3 = nn.Sequential(
+            nn.ConvTranspose2d(128, 64, 3, 2, 1),  # 38x38 -> 75x75
+            nn.LeakyReLU()
+        )
+
+        self.decoder4 = nn.Sequential(
+            nn.ConvTranspose2d(64, 3, 3, 2, 1, output_padding=1),  # 75x75 -> 150x150
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        # Encoder
+        enc1 = self.encoder1(x)  # 64 channels, 75x75
+        enc2 = self.encoder2(enc1)  # 128 channels, 38x38
+        enc3 = self.encoder3(enc2)  # 256 channels, 19x19
+        enc4 = self.encoder4(enc3)  # 512 channels, 10x10
+
+        # Bottleneck
+        bottleneck = self.bottleneck(enc4)
+
+        # Decoder (with skip connections)
+        dec_fc = self.decoder_fc(bottleneck)
+        dec1 = self.decoder1(dec_fc + enc4)  # Skip connection from encoder4
+        dec2 = self.decoder2(dec1 + enc3)  # Skip connection from encoder3
+        dec3 = self.decoder3(dec2 + enc2)  # Skip connection from encoder2
+        dec4 = self.decoder4(dec3 + enc1)  # Skip connection from encoder1
+
+        return dec4
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = DummyModel()
+model.load_state_dict(torch.load('model.pth'))
+model = model.to(device)
+model.eval()
+
+# Define preprocessing transforms
+preprocess = transforms.Compose([
+    transforms.Resize(150),
+    transforms.ToTensor(),
+    transforms.Normalize([0.5), [0.5])
+])
+
+def predict(image):
+    image = preprocess(image).to(model.device)
+    with torch.no_grad():
+        output = model(image)
+
+    image = transforms.ToPILImage()(output.squeeze().cpu())
+
+    return image
+
+# Create Gradio interface
+iface = gr.Interface(fn=predict,
+                     inputs=gr.Image(type="pil"),
+                     outputs=gr.Image(type="pil"))
+
+iface.launch()