Add Gradio app for ISIC skin lesion segmentation

README.md

@@ -1,12 +1,35 @@
 ---
-title: Isic Segmentation
-emoji: 🐠
-colorFrom: purple
-colorTo: green
+title: ISIC Skin Lesion Segmentation
+emoji: 🔬
+colorFrom: indigo
+colorTo: purple
 sdk: gradio
-sdk_version: 6.9.0
+sdk_version: "4.44.0"
 app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# ISIC 2018 Skin Lesion Segmentation — U-Net
+
+Upload a dermoscopy image to get an instant binary segmentation mask from a trained U-Net.
+
+## Results
+
+| Metric | Test Set |
+|--------|----------|
+| Dice   | **0.9301 ± 0.0621** |
+| IoU    | **0.8744 ± 0.0891** |
+
+Trained on **ISIC 2018 Task 1** (568 images, 70/15/15 train/val/test split).  
+Best checkpoint: epoch **45**, val Dice **0.9207**.
+
+## Model Architecture
+
+Classic **U-Net** with skip connections.  
+Channel progression: 3 → 64 → 128 → 256 → 512 → 1024 → 512 → 256 → 128 → 64 → 1
+
+## Usage
+
+1. Upload any dermoscopy / skin lesion image
+2. Click **Segment 🔍**
+3. View the predicted binary mask and overlay

app.py

@@ -0,0 +1,139 @@
+"""
+app.py
+------
+Gradio demo for ISIC 2018 Skin Lesion Segmentation using a trained U-Net.
+
+Hosted on Hugging Face Spaces.
+Model weights are downloaded from the HF Hub model repo on first run.
+"""
+
+import os
+import numpy as np
+import torch
+import gradio as gr
+from PIL import Image
+from huggingface_hub import hf_hub_download
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+MODEL_REPO   = "pavanpraneeth/isic-unet"
+MODEL_FILE   = "best_model.pth"
+IMAGE_SIZE   = 256
+THRESHOLD    = 0.5
+IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+IMAGENET_STD  = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+
+DEVICE = (
+    torch.device("cuda") if torch.cuda.is_available()
+    else torch.device("mps") if torch.backends.mps.is_available()
+    else torch.device("cpu")
+)
+
+# ---------------------------------------------------------------------------
+# Load model (once at startup)
+# ---------------------------------------------------------------------------
+from model import UNet  # model.py is alongside app.py in the Space repo
+
+def load_model() -> torch.nn.Module:
+    ckpt_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
+    model = UNet(in_channels=3, out_channels=1)
+    state = torch.load(ckpt_path, map_location=DEVICE)
+    model.load_state_dict(state["model_state_dict"])
+    model.eval().to(DEVICE)
+    print(f"[app] Model loaded from {MODEL_REPO} on {DEVICE}")
+    return model
+
+MODEL = load_model()
+
+# ---------------------------------------------------------------------------
+# Preprocessing / postprocessing helpers
+# ---------------------------------------------------------------------------
+
+def preprocess(img: np.ndarray) -> torch.Tensor:
+    """Resize, normalise (ImageNet), convert to tensor."""
+    pil = Image.fromarray(img).convert("RGB").resize((IMAGE_SIZE, IMAGE_SIZE))
+    arr = np.array(pil, dtype=np.float32) / 255.0
+    arr = (arr - IMAGENET_MEAN) / IMAGENET_STD          # (H, W, 3)
+    tensor = torch.from_numpy(arr.transpose(2, 0, 1))   # (3, H, W)
+    return tensor.unsqueeze(0).to(DEVICE)               # (1, 3, H, W)
+
+
+def postprocess_mask(pred: torch.Tensor) -> np.ndarray:
+    """Convert raw sigmoid output → uint8 mask image (0 or 255)."""
+    mask = (pred.squeeze().cpu().numpy() > THRESHOLD).astype(np.uint8) * 255
+    return mask
+
+
+def make_overlay(original_rgb: np.ndarray, mask: np.ndarray) -> np.ndarray:
+    """Overlay mask boundary on original image in red."""
+    h, w = mask.shape
+    orig_resized = np.array(
+        Image.fromarray(original_rgb).resize((w, h))
+    ).copy()
+
+    # Draw red where mask == 255
+    overlay = orig_resized.copy()
+    overlay[mask > 0] = (
+        overlay[mask > 0] * 0.4 + np.array([255, 0, 0]) * 0.6
+    ).astype(np.uint8)
+    return overlay
+
+
+# ---------------------------------------------------------------------------
+# Inference function (called by Gradio)
+# ---------------------------------------------------------------------------
+
+def segment(image: np.ndarray):
+    """Run inference and return (mask_image, overlay_image)."""
+    if image is None:
+        return None, None
+
+    tensor = preprocess(image)
+    with torch.no_grad():
+        pred = MODEL(tensor)         # (1, 1, 256, 256)
+
+    mask   = postprocess_mask(pred)  # (256, 256) uint8
+    overlay = make_overlay(image, mask)
+
+    mask_rgb = np.stack([mask, mask, mask], axis=-1)  # grey → RGB for display
+    return mask_rgb, overlay
+
+
+# ---------------------------------------------------------------------------
+# Gradio UI
+# ---------------------------------------------------------------------------
+
+DESCRIPTION = """
+## 🔬 ISIC 2018 Skin Lesion Segmentation
+
+Upload a dermoscopy image to get an instant binary segmentation mask from a trained **U-Net**.
+
+| Metric | Test Set Score |
+|--------|---------------|
+| Dice   | **0.9301 ± 0.0621** |
+| IoU    | **0.8744 ± 0.0891** |
+
+*Trained on ISIC 2018 Task 1 (568 images, 70/15/15 split).*
+"""
+
+with gr.Blocks(theme=gr.themes.Soft(), title="ISIC Skin Lesion Segmentation") as demo:
+    gr.Markdown(DESCRIPTION)
+
+    with gr.Row():
+        with gr.Column():
+            inp = gr.Image(label="Input Image", type="numpy")
+            btn = gr.Button("Segment 🔍", variant="primary")
+        with gr.Column():
+            out_mask    = gr.Image(label="Predicted Mask")
+            out_overlay = gr.Image(label="Overlay on Original")
+
+    btn.click(fn=segment, inputs=inp, outputs=[out_mask, out_overlay])
+
+    gr.Examples(
+        examples=[],
+        inputs=inp,
+    )
+
+if __name__ == "__main__":
+    demo.launch()

model.py

@@ -0,0 +1,74 @@
+"""
+model.py  —  Self-contained UNet definition for HF Spaces.
+Copied from the training project (model.py).
+"""
+
+import torch
+import torch.nn as nn
+from typing import List
+
+
+class DoubleConv(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int) -> None:
+        super().__init__()
+        self.block = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.block(x)
+
+
+class UNet(nn.Module):
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 1,
+        features: List[int] = [64, 128, 256, 512],
+    ) -> None:
+        super().__init__()
+
+        self.encoders = nn.ModuleList()
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+        ch = in_channels
+        for feat in features:
+            self.encoders.append(DoubleConv(ch, feat))
+            ch = feat
+
+        self.bottleneck = DoubleConv(features[-1], features[-1] * 2)
+        ch = features[-1] * 2
+
+        self.up_convs = nn.ModuleList()
+        self.decoders = nn.ModuleList()
+        for feat in reversed(features):
+            self.up_convs.append(nn.ConvTranspose2d(ch, feat, kernel_size=2, stride=2))
+            self.decoders.append(DoubleConv(feat * 2, feat))
+            ch = feat
+
+        self.output_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        skip_connections: List[torch.Tensor] = []
+
+        for encoder in self.encoders:
+            x = encoder(x)
+            skip_connections.append(x)
+            x = self.pool(x)
+
+        x = self.bottleneck(x)
+
+        skip_connections = skip_connections[::-1]
+        for up_conv, decoder, skip in zip(self.up_convs, self.decoders, skip_connections):
+            x = up_conv(x)
+            if x.shape != skip.shape:
+                x = nn.functional.interpolate(x, size=skip.shape[2:], mode="bilinear", align_corners=False)
+            x = torch.cat([skip, x], dim=1)
+            x = decoder(x)
+
+        return self.sigmoid(self.output_conv(x))

requirements.txt

@@ -0,0 +1,6 @@
+torch==2.2.2
+torchvision==0.17.2
+gradio==4.44.0
+huggingface_hub==0.23.4
+Pillow
+numpy