Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,504 @@
+"""
+Brain Tumor Detection — Gradio Space
+Hybrid CNN-ViT model with Grad-CAM explainability.
+
+Author: Vishnu K (ZorroJurro)
+"""
+
+import os
+import json
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+from torchvision import transforms
+from PIL import Image
+import cv2
+import gradio as gr
+from huggingface_hub import hf_hub_download
+from einops import rearrange, repeat
+from typing import Dict, Optional, Tuple
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+
+# =============================================================================
+# Model Architecture (self-contained)
+# =============================================================================
+
+class CNNBackbone(nn.Module):
+    def __init__(self, backbone_name="resnet50", pretrained=False, output_features=True):
+        super().__init__()
+        self.backbone_name = backbone_name.lower()
+        self.output_features = output_features
+        configs = {
+            "resnet50": (models.resnet50, models.ResNet50_Weights.IMAGENET1K_V2, 2048),
+        }
+        model_fn, weights, self.num_features = configs[self.backbone_name]
+        model = model_fn(weights=weights if pretrained else None)
+        self.backbone = nn.Sequential(*list(model.children())[:-2])
+
+    def forward(self, x):
+        return self.backbone(x)
+
+
+class PatchEmbedding(nn.Module):
+    def __init__(self, feature_size=7, feature_dim=2048, embed_dim=512, patch_size=1):
+        super().__init__()
+        self.patch_size = patch_size
+        self.num_patches = (feature_size // patch_size) ** 2
+        self.projection = nn.Linear(feature_dim, embed_dim) if patch_size == 1 else nn.Conv2d(feature_dim, embed_dim, kernel_size=patch_size, stride=patch_size)
+        self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim) * 0.02)
+        self.pos_embedding = nn.Parameter(torch.randn(1, self.num_patches + 1, embed_dim) * 0.02)
+
+    def forward(self, x):
+        B = x.shape[0]
+        if self.patch_size == 1:
+            x = rearrange(x, "b c h w -> b (h w) c")
+            x = self.projection(x)
+        else:
+            x = self.projection(x)
+            x = rearrange(x, "b c h w -> b (h w) c")
+        cls_tokens = repeat(self.cls_token, "1 1 d -> b 1 d", b=B)
+        x = torch.cat([cls_tokens, x], dim=1)
+        x = x + self.pos_embedding[:, :x.size(1)]
+        return x
+
+
+class MultiHeadSelfAttention(nn.Module):
+    def __init__(self, embed_dim=512, num_heads=8, dropout=0.1, attention_dropout=0.1):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        self.scale = self.head_dim ** -0.5
+        self.qkv = nn.Linear(embed_dim, embed_dim * 3)
+        self.attn_dropout = nn.Dropout(attention_dropout)
+        self.proj = nn.Linear(embed_dim, embed_dim)
+        self.proj_dropout = nn.Dropout(dropout)
+        self.attention_weights = None
+
+    def forward(self, x, return_attention=False):
+        B, N, D = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_dropout(attn)
+        self.attention_weights = attn.detach()
+        x = (attn @ v).transpose(1, 2).reshape(B, N, D)
+        x = self.proj_dropout(self.proj(x))
+        return (x, attn) if return_attention else (x, None)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, embed_dim=512, num_heads=8, mlp_ratio=4.0, dropout=0.1, attention_dropout=0.1):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(embed_dim)
+        self.attn = MultiHeadSelfAttention(embed_dim, num_heads, dropout, attention_dropout)
+        self.norm2 = nn.LayerNorm(embed_dim)
+        mlp_hidden = int(embed_dim * mlp_ratio)
+        self.mlp = nn.Sequential(nn.Linear(embed_dim, mlp_hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(mlp_hidden, embed_dim), nn.Dropout(dropout))
+        self.attention_weights = None
+
+    def forward(self, x, return_attention=False):
+        attn_out, attn = self.attn(self.norm1(x), return_attention)
+        x = x + attn_out
+        x = x + self.mlp(self.norm2(x))
+        self.attention_weights = attn
+        return (x, attn) if return_attention else (x, None)
+
+
+class ViTEncoder(nn.Module):
+    def __init__(self, embed_dim=512, depth=6, num_heads=8, mlp_ratio=4.0, dropout=0.1, attention_dropout=0.1):
+        super().__init__()
+        self.blocks = nn.ModuleList([TransformerBlock(embed_dim, num_heads, mlp_ratio, dropout, attention_dropout) for _ in range(depth)])
+        self.norm = nn.LayerNorm(embed_dim)
+        self.attention_weights_all = []
+
+    def forward(self, x, return_attention=False):
+        self.attention_weights_all = []
+        for block in self.blocks:
+            x, attn = block(x, return_attention)
+            if return_attention and attn is not None:
+                self.attention_weights_all.append(attn)
+        x = self.norm(x)
+        return (x, self.attention_weights_all) if return_attention else (x, None)
+
+
+class LearnableRadiomics(nn.Module):
+    def __init__(self, in_channels=3, feature_dim=128):
+        super().__init__()
+        self.texture_branch = nn.Sequential(nn.Conv2d(in_channels, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(), nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(), nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(64, feature_dim // 2))
+        self.shape_branch = nn.Sequential(nn.Conv2d(in_channels, 32, 5, padding=2), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, 5, padding=2), nn.BatchNorm2d(64), nn.ReLU(), nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(64, feature_dim // 2))
+        self.fusion = nn.Sequential(nn.Linear(feature_dim, feature_dim), nn.LayerNorm(feature_dim), nn.ReLU())
+
+    def forward(self, x):
+        return self.fusion(torch.cat([self.texture_branch(x), self.shape_branch(x)], dim=-1))
+
+
+class FeatureFusion(nn.Module):
+    def __init__(self, cnn_dim=2048, vit_dim=512, radiomics_dim=128, output_dim=512, use_radiomics=True):
+        super().__init__()
+        self.use_radiomics = use_radiomics
+        total_dim = cnn_dim + vit_dim + (radiomics_dim if use_radiomics else 0)
+        self.fusion = nn.Sequential(nn.Linear(total_dim, output_dim * 2), nn.LayerNorm(output_dim * 2), nn.GELU(), nn.Dropout(0.1), nn.Linear(output_dim * 2, output_dim), nn.LayerNorm(output_dim), nn.GELU())
+
+    def forward(self, cnn_features, vit_features, radiomics_features=None):
+        parts = [cnn_features, vit_features]
+        if self.use_radiomics and radiomics_features is not None:
+            parts.append(radiomics_features)
+        return self.fusion(torch.cat(parts, dim=-1))
+
+
+class HybridCNNViT(nn.Module):
+    def __init__(self, num_classes=4, cnn_backbone="resnet50", cnn_pretrained=False,
+                 vit_embed_dim=512, vit_depth=6, vit_num_heads=8, vit_mlp_ratio=4.0,
+                 use_radiomics=True, radiomics_dim=128, fusion_type="concat", dropout=0.3):
+        super().__init__()
+        self.use_radiomics = use_radiomics
+        self.cnn = CNNBackbone(backbone_name=cnn_backbone, pretrained=cnn_pretrained, output_features=True)
+        cnn_feature_dim = self.cnn.num_features
+        self.patch_embed = PatchEmbedding(feature_size=7, feature_dim=cnn_feature_dim, embed_dim=vit_embed_dim, patch_size=1)
+        self.vit_encoder = ViTEncoder(embed_dim=vit_embed_dim, depth=vit_depth, num_heads=vit_num_heads, mlp_ratio=vit_mlp_ratio, dropout=dropout * 0.5)
+        self.cnn_pool = nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Flatten())
+        self.radiomics = LearnableRadiomics(in_channels=3, feature_dim=radiomics_dim) if use_radiomics else None
+        if not use_radiomics:
+            radiomics_dim = 0
+        self.fusion = FeatureFusion(cnn_dim=cnn_feature_dim, vit_dim=vit_embed_dim, radiomics_dim=radiomics_dim, output_dim=512, use_radiomics=use_radiomics)
+        self.classifier = nn.Sequential(nn.Dropout(dropout), nn.Linear(512, 256), nn.LayerNorm(256), nn.GELU(), nn.Dropout(dropout * 0.5), nn.Linear(256, num_classes))
+        self.attention_weights = None
+
+    def forward(self, x, return_features=False, return_attention=False):
+        cnn_features = self.cnn(x)
+        cnn_pooled = self.cnn_pool(cnn_features)
+        patch_embeddings = self.patch_embed(cnn_features)
+        vit_output, attention = self.vit_encoder(patch_embeddings, return_attention)
+        vit_cls = vit_output[:, 0]
+        if return_attention:
+            self.attention_weights = attention
+        radiomics_features = self.radiomics(x) if self.use_radiomics else None
+        fused = self.fusion(cnn_pooled, vit_cls, radiomics_features)
+        logits = self.classifier(fused)
+        output = {"logits": logits}
+        if return_features:
+            output["cnn_features"] = cnn_pooled
+            output["vit_features"] = vit_cls
+            output["fused_features"] = fused
+        if return_attention:
+            output["attention"] = attention
+        return output
+
+
+class BrainTumorClassifier(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        mc = config.get("model", {})
+        self.model = HybridCNNViT(
+            num_classes=config.get("data", {}).get("num_classes", 4),
+            cnn_backbone=mc.get("cnn_backbone", "resnet50"),
+            cnn_pretrained=mc.get("cnn_pretrained", False),
+            vit_embed_dim=mc.get("vit_embed_dim", 512),
+            vit_depth=mc.get("vit_depth", 6),
+            vit_num_heads=mc.get("vit_num_heads", 8),
+            vit_mlp_ratio=mc.get("vit_mlp_ratio", 4.0),
+            use_radiomics=mc.get("use_radiomics", True),
+            radiomics_dim=mc.get("radiomics_features", 128),
+            dropout=mc.get("dropout", 0.3),
+        )
+        self.num_classes = config.get("data", {}).get("num_classes", 4)
+
+    def forward(self, x):
+        return self.model(x)["logits"]
+
+
+# =============================================================================
+# Grad-CAM Implementation
+# =============================================================================
+
+class GradCAM:
+    """Simplified Grad-CAM for the CNN backbone."""
+
+    def __init__(self, model: HybridCNNViT):
+        self.model = model
+        self.gradients = None
+        self.activations = None
+        self._register_hooks()
+
+    def _register_hooks(self):
+        # Hook into the last conv layer of the CNN backbone
+        target_layer = self.model.cnn.backbone[-1]
+
+        def forward_hook(module, input, output):
+            self.activations = output.detach()
+
+        def backward_hook(module, grad_input, grad_output):
+            self.gradients = grad_output[0].detach()
+
+        target_layer.register_forward_hook(forward_hook)
+        target_layer.register_full_backward_hook(backward_hook)
+
+    def generate(self, input_tensor: torch.Tensor, target_class: int = None) -> np.ndarray:
+        self.model.eval()
+        input_tensor.requires_grad_(True)
+
+        output = self.model(input_tensor)
+        logits = output["logits"]
+
+        if target_class is None:
+            target_class = logits.argmax(dim=-1).item()
+
+        self.model.zero_grad()
+        logits[0, target_class].backward()
+
+        gradients = self.gradients
+        activations = self.activations
+
+        # Global average pooling of gradients
+        weights = gradients.mean(dim=(2, 3), keepdim=True)
+        cam = (weights * activations).sum(dim=1, keepdim=True)
+        cam = F.relu(cam)
+
+        # Normalize
+        cam = cam.squeeze().cpu().numpy()
+        cam = cam - cam.min()
+        cam = cam / (cam.max() + 1e-8)
+
+        return cam
+
+
+def create_gradcam_overlay(
+    original_image: np.ndarray,
+    cam: np.ndarray,
+    alpha: float = 0.5,
+) -> np.ndarray:
+    """Create a Grad-CAM heatmap overlay on the original image."""
+    h, w = original_image.shape[:2]
+    cam_resized = cv2.resize(cam, (w, h))
+
+    # Apply colormap
+    heatmap = cv2.applyColorMap(np.uint8(255 * cam_resized), cv2.COLORMAP_JET)
+    heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
+
+    # Overlay
+    overlay = np.float32(heatmap) * alpha + np.float32(original_image) * (1 - alpha)
+    overlay = np.clip(overlay, 0, 255).astype(np.uint8)
+
+    return overlay
+
+
+# =============================================================================
+# Model Loading
+# =============================================================================
+
+REPO_ID = "Zorrojurro/brain-tumor-cnn-vit"
+CLASS_NAMES = ["Glioma", "Meningioma", "No Tumor", "Pituitary"]
+CLASS_EMOJIS = {"Glioma": "🔴", "Meningioma": "🟠", "No Tumor": "🟢", "Pituitary": "🟡"}
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Image preprocessing
+TRANSFORM = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+])
+
+
+def load_model():
+    """Download and load the model from Hugging Face Hub."""
+    print("📥 Downloading model from Hugging Face Hub...")
+
+    # Download checkpoint
+    checkpoint_path = hf_hub_download(
+        repo_id=REPO_ID,
+        filename="best_model.pth",
+        cache_dir="./model_cache",
+    )
+
+    # Create model
+    config = {
+        "data": {"num_classes": 4},
+        "model": {
+            "cnn_backbone": "resnet50",
+            "cnn_pretrained": False,
+            "vit_embed_dim": 512,
+            "vit_depth": 6,
+            "vit_num_heads": 8,
+            "vit_mlp_ratio": 4.0,
+            "use_radiomics": True,
+            "radiomics_features": 128,
+            "dropout": 0.3,
+        },
+    }
+
+    classifier = BrainTumorClassifier(config)
+    model = classifier.model
+
+    # Load checkpoint
+    checkpoint = torch.load(checkpoint_path, map_location=DEVICE, weights_only=False)
+    state_dict = checkpoint.get("model_state_dict", checkpoint)
+
+    # Handle key prefix mismatches
+    new_state_dict = {}
+    for k, v in state_dict.items():
+        # Remove 'model.' prefix if present
+        new_key = k.replace("model.", "") if k.startswith("model.") else k
+        new_state_dict[new_key] = v
+
+    model.load_state_dict(new_state_dict, strict=False)
+    model.eval().to(DEVICE)
+
+    print(f"✅ Model loaded on {DEVICE}")
+    return model
+
+
+# Load model at startup
+MODEL = load_model()
+GRADCAM = GradCAM(MODEL)
+
+
+# =============================================================================
+# Prediction Function
+# =============================================================================
+
+def predict(image: Image.Image):
+    """Run prediction and generate Grad-CAM visualization."""
+    if image is None:
+        return None, None, "Please upload an image."
+
+    # Convert to RGB
+    image = image.convert("RGB")
+    original_np = np.array(image)
+
+    # Preprocess
+    input_tensor = TRANSFORM(image).unsqueeze(0).to(DEVICE)
+
+    # Forward pass with gradients for Grad-CAM
+    with torch.enable_grad():
+        cam = GRADCAM.generate(input_tensor)
+
+    # Get predictions
+    with torch.no_grad():
+        output = MODEL(input_tensor)
+        logits = output["logits"]
+        probs = F.softmax(logits, dim=-1)[0]
+
+    # Build confidence dict
+    confidences = {}
+    for i, name in enumerate(CLASS_NAMES):
+        emoji = CLASS_EMOJIS[name]
+        confidences[f"{emoji} {name}"] = float(probs[i])
+
+    # Grad-CAM overlay
+    gradcam_overlay = create_gradcam_overlay(original_np, cam, alpha=0.45)
+
+    # Predicted class info
+    pred_idx = probs.argmax().item()
+    pred_name = CLASS_NAMES[pred_idx]
+    pred_conf = probs[pred_idx].item()
+    emoji = CLASS_EMOJIS[pred_name]
+
+    summary = f"## {emoji} {pred_name}\n**Confidence:** {pred_conf:.1%}\n\n"
+    if pred_name == "No Tumor":
+        summary += "✅ No tumor detected in the MRI scan."
+    else:
+        summary += f"⚠️ Potential **{pred_name.lower()}** detected. Please consult a medical professional."
+
+    return confidences, gradcam_overlay, summary
+
+
+# =============================================================================
+# Gradio UI
+# =============================================================================
+
+CUSTOM_CSS = """
+.gradio-container {
+    max-width: 1100px !important;
+    margin: auto !important;
+}
+.gr-button-primary {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
+    border: none !important;
+}
+.gr-button-primary:hover {
+    background: linear-gradient(135deg, #764ba2 0%, #667eea 100%) !important;
+    transform: translateY(-1px);
+    box-shadow: 0 4px 15px rgba(102, 126, 234, 0.4);
+}
+footer {visibility: hidden}
+"""
+
+DESCRIPTION = """
+# 🧠 Brain Tumor Detection — Hybrid CNN-ViT
+
+Upload a brain MRI scan for instant AI-powered classification with **Grad-CAM explainability**.
+
+**Model Architecture**: ResNet50 (CNN) + 6-Layer Vision Transformer + Learnable Radiomics  
+**Classes**: Glioma · Meningioma · No Tumor · Pituitary  
+**Performance**: 98% Accuracy · 0.97 F1-Score · 0.99 AUC
+
+> ⚠️ *For research and educational purposes only. Not a substitute for professional medical diagnosis.*
+"""
+
+with gr.Blocks(
+    css=CUSTOM_CSS,
+    theme=gr.themes.Soft(
+        primary_hue="indigo",
+        secondary_hue="purple",
+        neutral_hue="slate",
+    ),
+    title="Brain Tumor Detection — CNN-ViT",
+) as demo:
+
+    gr.Markdown(DESCRIPTION)
+
+    with gr.Row(equal_height=True):
+        with gr.Column(scale=1):
+            input_image = gr.Image(
+                type="pil",
+                label="Upload Brain MRI",
+                height=350,
+            )
+            predict_btn = gr.Button(
+                "🔬 Analyze MRI",
+                variant="primary",
+                size="lg",
+            )
+
+        with gr.Column(scale=1):
+            gradcam_output = gr.Image(
+                label="Grad-CAM Visualization",
+                height=350,
+            )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            label_output = gr.Label(
+                label="Classification Confidence",
+                num_top_classes=4,
+            )
+        with gr.Column(scale=1):
+            summary_output = gr.Markdown(
+                label="Diagnosis Summary",
+            )
+
+    predict_btn.click(
+        fn=predict,
+        inputs=[input_image],
+        outputs=[label_output, gradcam_output, summary_output],
+    )
+
+    gr.Markdown(
+        """
+        ---
+        **Built by [Vishnu K](https://huggingface.co/ZorroJurro)** · 
+        [Model Card](https://huggingface.co/ZorroJurro/brain-tumor-cnn-vit) · 
+        [GitHub](https://github.com/ZorroJurro)
+        """
+    )
+
+
+if __name__ == "__main__":
+    demo.launch()