Update app.py

app.py

@@ -6,11 +6,19 @@ import gradio as gr
 from huggingface_hub import hf_hub_download
 import os
 
+# ===============================
+# 1. LOAD MODEL FROM HF
+# ===============================
+
 model_path = hf_hub_download(
     repo_id="harikrishnaaa321/cnn_attention_model",
     filename="cnn_attention_best.pth"
 )
 
+# ===============================
+# 2. MODEL ARCHITECTURE
+# ===============================
+
 class SEBlock(nn.Module):
     def __init__(self, channels, reduction=8):
         super(SEBlock, self).__init__()
@@ -57,6 +65,10 @@ class CNN_Attention_Model(nn.Module):
         return x
 
 
+# ===============================
+# 3. LOAD MODEL
+# ===============================
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = CNN_Attention_Model(num_classes=4).to(device)
 state_dict = torch.load(model_path, map_location=device)
@@ -64,24 +76,34 @@ model.load_state_dict(state_dict, strict=False)
 model.eval()
 
 labels = ["Glioma", "Meningioma", "Pituitary", "Normal"]
-EXAMPLE_IMAGE_DIR = "./examples"  # Must contain glioma.jpg, meningioma.jpg, pituitary.jpg, notumor.jpg
+
+# Direct images (not folder)
 example_images = {
-    "Glioma": os.path.join(EXAMPLE_IMAGE_DIR, "glioma.jpg"),
-    "Meningioma": os.path.join(EXAMPLE_IMAGE_DIR, "meningioma.jpg"),
-    "Pituitary": os.path.join(EXAMPLE_IMAGE_DIR, "pituitary.jpg"),
-    "Normal": os.path.join(EXAMPLE_IMAGE_DIR, "notumor.jpg")
+    "Glioma": "./glioma.jpg",
+    "Meningioma": "./meningioma.jpg",
+    "Pituitary": "./pituitary.jpg",
+    "Normal": "./notumor.jpg"
 }
 
+# ===============================
+# 4. TRANSFORMS
+# ===============================
+
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
 ])
 
+
+# ===============================
+# 5. PREDICT FUNCTION
+# ===============================
+
 def predict_tumor(image):
     # Convert to grayscale
     image = image.convert("L")
 
-    # Transform to tensor [1, 2, H, W]
+    # Make 2-channel input
     img_tensor = transform(image).repeat(2, 1, 1).unsqueeze(0).to(device)
 
     # Forward pass
@@ -93,30 +115,64 @@ def predict_tumor(image):
     pred_label = labels[pred_idx]
     confidences = {labels[i]: float(probs[i]) for i in range(len(labels))}
 
-    # Create combined image with names below (dark background)
+    # ===============================
+    # Create reference panel (CREAM BG)
+    # ===============================
+
     example_imgs = [Image.open(example_images[label]).resize((224, 224)) for label in labels]
+
     combined_width = 224 * 4
-    combined_height = 224 + 30
-    combined_image = Image.new("RGB", (combined_width, combined_height), "#1a1a1a")  # Dark background
+    combined_height = 224 + 40
+
+    combined_image = Image.new("RGB", (combined_width, combined_height), "#f8eecf")  # cream
+
     draw = ImageDraw.Draw(combined_image)
     font = ImageFont.load_default()
 
     for i, img in enumerate(example_imgs):
-        combined_image.paste(img, (224 * i, 0))
-        # Draw red border for prediction
+        x = 224 * i
+        combined_image.paste(img, (x, 0))
+
+        # Highlight predicted class with red border
         if i == pred_idx:
-            draw.rectangle([224 * i, 0, 224 * (i + 1) - 1, 223], outline="#ff4444", width=5)
-        # Draw label text below in white
+            draw.rectangle([x, 0, x + 223, 223], outline="#ff4444", width=5)
+
+        # Draw label text
         text = labels[i]
         bbox = draw.textbbox((0, 0), text, font=font)
-        text_width = bbox[2] - bbox[0]
-        text_height = bbox[3] - bbox[1]
-        text_x = 224 * i + (224 - text_width) // 2
-        text_y = 224 + 5
-        draw.text((text_x, text_y), text, fill="white", font=font)  # White text for dark theme
+        tw = bbox[2] - bbox[0]
+        tx = x + (224 - tw) // 2
+        draw.text((tx, 228), text, fill="black", font=font)
 
     return pred_label, confidences, combined_image
 
+
+# ===============================
+# 6. CREAM UI THEME CSS
+# ===============================
+
+custom_css = """
+:root {
+    --body-background-fill: #f8eecf;
+    --block-background-fill: #fffaf0;
+    --border-color: #d5c7a1;
+    --button-primary-background-fill: #d6b77a;
+    --button-primary-text-color: #000000;
+}
+
+body, .gradio-container {
+    background: #f8eecf !important;
+}
+
+.gr-button {
+    border-radius: 8px !important;
+}
+"""
+
+# ===============================
+# 7. GRADIO INTERFACE
+# ===============================
+
 interface = gr.Interface(
     fn=predict_tumor,
     inputs=gr.Image(type="pil", label="Upload MRI Scan"),
@@ -125,34 +181,10 @@ interface = gr.Interface(
         gr.Label(label="Confidence Scores"),
         gr.Image(label="Reference Tumor Images")
     ],
-    title="🧠 Brain Tumor Classification",
-    description="Upload an MRI scan (JPG/PNG). The model will predict the tumor type and show reference images with names.",
-    theme=gr.themes.Soft(
-        primary_hue="blue",
-        secondary_hue="slate",
-    ).set(
-        body_background_fill="#0f0f0f",
-        body_background_fill_dark="#0f0f0f",
-        block_background_fill="#1a1a1a",
-        block_background_fill_dark="#1a1a1a",
-        block_border_width="1px",
-        block_label_background_fill="#1a1a1a",
-        block_label_background_fill_dark="#1a1a1a",
-        block_label_text_color="#60a5fa",
-        block_label_text_color_dark="#60a5fa",
-        block_title_text_color="#3b82f6",
-        block_title_text_color_dark="#3b82f6",
-        body_text_color="#e5e7eb",
-        body_text_color_dark="#e5e7eb",
-        input_background_fill="#2d2d2d",
-        input_background_fill_dark="#2d2d2d",
-        button_primary_background_fill="#3b82f6",
-        button_primary_background_fill_hover="#2563eb",
-        # Label text inside components (like confidence scores)
-        color_accent="#60a5fa",
-        color_accent_soft="#3b82f6",
-    ),
+    title="🧠 Brain Tumor Classification (Attention CNN)",
+    description="Upload an MRI scan. Model predicts tumor type and shows reference examples.",
+    css=custom_css
 )
 
 if __name__ == "__main__":
-    interface.launch()
+    interface.launch()