Upload 3 files

app.py

@@ -0,0 +1,174 @@
+import torch
+import torch.nn as nn
+import torchvision.transforms as transforms
+import gradio as gr
+import numpy as np
+from PIL import Image
+import cv2
+
+############################################
+# ========== UNET MODEL ====================
+############################################
+
+class DoubleConv(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.conv_op = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.conv_op(x)
+
+class DownSample(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.conv = DoubleConv(in_channels, out_channels)
+        self.pool = nn.MaxPool2d(2)
+
+    def forward(self, x):
+        down = self.conv(x)
+        p = self.pool(down)
+        return down, p
+
+class UpSample(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.up = nn.ConvTranspose2d(in_channels, in_channels//2, 2, 2)
+        self.conv = DoubleConv(in_channels, out_channels)
+
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        x = torch.cat([x1, x2], 1)
+        return self.conv(x)
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=3, num_classes=1):
+        super().__init__()
+
+        self.down1 = DownSample(in_channels, 64)
+        self.down2 = DownSample(64, 128)
+        self.down3 = DownSample(128, 256)
+        self.down4 = DownSample(256, 512)
+
+        self.bottleneck = DoubleConv(512, 1024)
+
+        self.up1 = UpSample(1024, 512)
+        self.up2 = UpSample(512, 256)
+        self.up3 = UpSample(256, 128)
+        self.up4 = UpSample(128, 64)
+
+        self.out = nn.Conv2d(64, num_classes, kernel_size=1)
+
+    def forward(self, x):
+        d1, p1 = self.down1(x)
+        d2, p2 = self.down2(p1)
+        d3, p3 = self.down3(p2)
+        d4, p4 = self.down4(p3)
+
+        b = self.bottleneck(p4)
+
+        u1 = self.up1(b, d4)
+        u2 = self.up2(u1, d3)
+        u3 = self.up3(u2, d2)
+        u4 = self.up4(u3, d1)
+
+        return self.out(u4)
+
+############################################
+# ========== LOAD MODEL ====================
+############################################
+
+device = torch.device("cpu")
+
+model = UNet()
+model.load_state_dict(torch.load("my_checkpoint.pth", map_location=device))
+model.eval()
+
+############################################
+# ========== TRANSFORM =====================
+############################################
+
+transform = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.ToTensor()
+])
+
+############################################
+# ========== DICE ==========================
+############################################
+
+def dice_coefficient(pred, target, epsilon=1e-7):
+    pred = (pred > 0.5).float()
+    intersection = (pred * target).sum()
+    union = pred.sum() + target.sum()
+    return ((2. * intersection + epsilon) / (union + epsilon)).item()
+
+############################################
+# ========== INFERENCE FUNCTION ============
+############################################
+
+def predict(image, mask=None):
+
+    image_pil = Image.fromarray(image).convert("RGB")
+    input_tensor = transform(image_pil).unsqueeze(0)
+
+    with torch.no_grad():
+        output = model(input_tensor)
+        output = torch.sigmoid(output)
+
+    pred_mask = output.squeeze().numpy()
+    pred_mask_binary = (pred_mask > 0.5).astype(np.uint8)
+
+    # Resize mask back to original size
+    pred_mask_resized = cv2.resize(
+        pred_mask_binary,
+        (image.shape[1], image.shape[0])
+    )
+
+    # Create overlay
+    overlay = image.copy()
+    overlay[pred_mask_resized == 1] = [255, 0, 0]
+
+    if mask is not None:
+        mask_pil = Image.fromarray(mask).convert("L")
+        mask_tensor = transform(mask_pil)
+        dice = dice_coefficient(torch.tensor(pred_mask), mask_tensor)
+        return overlay, f"Dice Score: {round(dice, 4)}"
+
+    return overlay, "Mask predicted successfully!"
+
+############################################
+# ========== GRADIO UI =====================
+############################################
+
+description = """
+# 🧠 Brain Tumor Segmentation (UNet)
+
+This model was trained on:
+
+🔗 https://www.kaggle.com/datasets/mateuszbuda/lgg-mri-segmentation
+
+Upload an MRI image to see tumor segmentation.
+Optionally upload the true mask to compute Dice score.
+"""
+
+demo = gr.Interface(
+    fn=predict,
+    inputs=[
+        gr.Image(type="numpy", label="Upload MRI Image"),
+        gr.Image(type="numpy", label="Optional Ground Truth Mask")
+    ],
+    outputs=[
+        gr.Image(label="Predicted Overlay"),
+        gr.Textbox(label="Info")
+    ],
+    title="UNet Brain Tumor Segmentation",
+    description=description
+)
+
+if __name__ == "__main__":
+    demo.launch()

my_checkpoint.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:33292122b7a931f081e57e9aec82a391c92f3d6154bec3f9e61d5a2c82a0c9dd
+size 124146577

requirements.txt

@@ -0,0 +1,7 @@
+torch
+torchvision
+gradio
+opencv-python
+numpy
+pillow
+matplotlib