Update app.py

app.py

@@ -1,174 +1,201 @@
-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__":
+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 = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, 3, padding=1),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x):
+        return self.conv(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], dim=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 FUNCTION =================
+############################################
+
+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()
+
+############################################
+# ========== PREPROCESS TIFF ===============
+############################################
+
+def load_image(file):
+    img = Image.open(file.name)
+    img_np = np.array(img)
+
+    # Handle 16-bit TIFF
+    if img_np.dtype == np.uint16:
+        img_np = (img_np / 256).astype(np.uint8)
+
+    img_pil = Image.fromarray(img_np).convert("RGB")
+    return img_pil, img_np
+
+
+############################################
+# ========== PREDICTION ====================
+############################################
+
+def predict(image_file, mask_file=None):
+
+    if image_file is None:
+        return None, "Please upload an image."
+
+    image_pil, original_np = load_image(image_file)
+
+    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 to original image size
+    pred_mask_resized = cv2.resize(
+        pred_mask_binary,
+        (original_np.shape[1], original_np.shape[0])
+    )
+
+    # Create red overlay
+    overlay = original_np.copy()
+    overlay[pred_mask_resized == 1] = [255, 0, 0]
+
+    # If mask provided → compute Dice
+    if mask_file is not None:
+        mask_pil, _ = load_image(mask_file)
+        mask_tensor = transform(mask_pil.convert("L"))
+        dice = dice_coefficient(torch.tensor(pred_mask), mask_tensor)
+        return overlay, f"Dice Score: {round(dice,4)}"
+
+    return overlay, "Prediction complete."
+
+
+############################################
+# ========== GRADIO UI =====================
+############################################
+
+description = """
+# 🧠 Brain Tumor Segmentation using UNet
+
+This model was trained on:
+
+🔗 https://www.kaggle.com/datasets/mateuszbuda/lgg-mri-segmentation
+
+Upload a `.tif` MRI image to predict tumor segmentation.
+Optionally upload the true mask to compute Dice score.
+"""
+
+demo = gr.Interface(
+    fn=predict,
+    inputs=[
+        gr.File(file_types=[".tif", ".tiff", ".png", ".jpg"], label="Upload MRI Image"),
+        gr.File(file_types=[".tif", ".tiff"], label="Optional Ground Truth Mask")
+    ],
+    outputs=[
+        gr.Image(label="Predicted Overlay"),
+        gr.Textbox(label="Result")
+    ],
+    title="UNet Brain Tumor Segmentation",
+    description=description,
+    allow_flagging="never"
+)
+
+if __name__ == "__main__":
     demo.launch()