Update app.py

app.py

@@ -4,118 +4,108 @@ from PIL import Image, ImageDraw, ImageFont
 import torch
 import numpy as np
 
-# --- Model Setup (Unchanged) ---
 model_name = "Hemgg/brain-tumor-classification"
 processor = AutoImageProcessor.from_pretrained(model_name)
 model = AutoModelForImageClassification.from_pretrained(model_name)
 class_names = ["🧠 Glioma", "🎯 Meningioma", "✅ No Tumor", "⚡ Pituitary"]
 
-# --- Custom CSS for Advanced Look ---
 custom_css = """
-body, .gradio-container {
-    background-color: #0b0f19 !important;
-    color: #ffffff !important;
-}
-.md, .md p, .md h1, .md h2, .md h3 {
-    color: white !important;
-}
-#custom-card {
-    background: rgba(255, 255, 255, 0.05);
-    border-radius: 15px;
+.gradio-container {background-color: #050505 !important; color: #ffffff !important;}
+.md, .md p, .md h1, .md h2, .md h3, span, label {color: #ffffff !important;}
+#result-box {
+    background-color: #1a1a1a !important; 
+    border: 2px solid #3b82f6 !important; 
+    border-radius: 12px; 
     padding: 20px;
-    border: 1px solid rgba(255, 255, 255, 0.1);
+    color: #ffffff !important;
 }
 .gr-button-primary {
-    background: linear-gradient(90deg, #2563eb, #7c3aed) !important;
+    background: linear-gradient(135deg, #1e40af, #7e22ce) !important;
     border: none !important;
+    font-weight: bold !important;
 }
+#result-text * { color: #ffffff !important; }
 footer {display: none !important;}
 """
 
+def get_medical_info(tumor_type):
+    info = {
+        "🧠 Glioma": {
+            "desc": "Gliomas originate in the glial cells that support neurons. They can be fast-growing and may involve surrounding brain tissue.",
+            "next": "Urgent consultation with a neuro-oncologist. An MRI with contrast or a biopsy is typically the next diagnostic step."
+        },
+        "🎯 Meningioma": {
+            "desc": "These tumors arise from the meninges, the layers covering the brain. Most are slow-growing and benign but can cause pressure.",
+            "next": "Consult a neurosurgeon to evaluate the mass effect. Treatment ranges from 'watchful waiting' to surgical resection."
+        },
+        "⚡ Pituitary": {
+            "desc": "Pituitary adenomas occur in the master gland at the base of the brain. They often affect hormone regulation and vision.",
+            "next": "An endocrine workup (blood tests) and a visual field test are recommended to assess hormonal and optic nerve impact."
+        },
+        "✅ No Tumor": {
+            "desc": "The neural network did not detect significant signs of the three primary tumor types in this scan.",
+            "next": "If symptoms (headaches, seizures, vision loss) persist, please consult a neurologist for a comprehensive evaluation."
+        }
+    }
+    return info.get(tumor_type, {"desc": "", "next": ""})
+
 def classify_tumor(image):
     if image is None:
-        return "Please upload a brain MRI", None
+        return "<p style='color:white;'>Please upload a scan.</p>", None
     
     inputs = processor(images=image, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
         probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
-        pred_idx = probs.argmax(-1).item()
-        confidence = probs[0][pred_idx].item()
-        tumor_type = class_names[pred_idx]
+        idx = probs.argmax(-1).item()
+        conf = probs[0][idx].item()
+        name = class_names[idx]
+        med = get_medical_info(name)
 
-    overlay_img = create_overlay(image, pred_idx, confidence)
+    overlay = create_overlay(image, idx, conf)
     
-    # Advanced Result Formatting
-    result = f"""
-    <div style="background: rgba(255,255,255,0.1); padding: 20px; border-radius: 10px; border-left: 5px solid #3b82f6;">
-        <h2 style="margin:0; color: white;">Analysis: {tumor_type}</h2>
-        <h3 style="margin:5px 0; color: #93c5fd;">Confidence Score: {confidence:.1%}</h3>
-        <hr style="opacity: 0.2; margin: 15px 0;">
-        <p style="color: #cbd5e1;"><b>Probability Distribution:</b></p>
-        {"".join([f"<div style='margin-bottom:5px;'>{class_names[i]}: <span style='float:right;'>{probs[0][i]:.1%}</span></div>" for i in range(4)])}
+    html_res = f"""
+    <div id="result-box">
+        <h2 style="color: #60a5fa; margin-top:0;">{name} Detected</h2>
+        <p style="font-size: 1.2em;"><b>Confidence Score:</b> <span style="color: #4ade80;">{conf:.1%}</span></p>
+        <div style="margin: 15px 0; padding: 10px; background: #262626; border-radius: 8px;">
+            <p style="color: #93c5fd; margin-bottom: 5px;"><b>Clinical Overview:</b></p>
+            <p style="color: #ffffff;">{med['desc']}</p>
+        </div>
+        <div style="margin: 15px 0; padding: 10px; background: #1e3a8a; border-radius: 8px;">
+            <p style="color: #bfdbfe; margin-bottom: 5px;"><b>Recommended Next Steps:</b></p>
+            <p style="color: #ffffff;">{med['next']}</p>
+        </div>
+        <p style="font-size: 0.8em; color: #94a3b8; border-top: 1px solid #444; pt-10;">
+            Probabilities: {", ".join([f"{class_names[i]}: {probs[0][i]:.1%}" for i in range(4)])}
+        </p>
     </div>
     """
-    return result, overlay_img
+    return html_res, overlay
 
-def create_overlay(image, tumor_class, confidence):
+def create_overlay(image, idx, conf):
     overlay = image.copy().convert("RGBA")
     draw = ImageDraw.Draw(overlay)
-    w, h = overlay.size
-    
-    if tumor_class != 2:
-        radius = min(w, h) // 4
+    if idx != 2:
+        w, h = overlay.size
+        r = min(w, h) // 4
         cx, cy = w // 2, h // 2
-        alpha = int(140 * confidence)
-        draw.ellipse([cx-radius, cy-radius, cx+radius, cy+radius], 
-                     fill=(255, 0, 0, alpha), outline=(255, 255, 0, 255), width=4)
-        
-        try:
-            font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 24)
-        except:
-            font = ImageFont.load_default()
-            
-        label = f"REGION OF INTEREST: {class_names[tumor_class]}"
-        draw.text((20, 20), label, fill=(255, 255, 255), font=font)
-            
+        draw.ellipse([cx-r, cy-r, cx+r, cy+r], fill=(255, 0, 0, int(150*conf)), outline=(255, 255, 0, 255), width=5)
     return overlay.convert("RGB")
 
-# --- Interface ---
-with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
+with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
+    gr.HTML("<h1 style='text-align:center; color:white; font-size: 32px;'>NEURO-DIAGNOSTIC AI STATION</h1>")
+    
     with gr.Row():
-        gr.HTML("""
-            <div style="text-align: center; padding: 20px;">
-                <h1 style="color: white; font-size: 2.5rem; margin-bottom: 0;">NEURO-SCAN AI</h1>
-                <p style="color: #94a3b8; font-size: 1.1rem;">Advanced Deep Learning Diagnostic Support System</p>
-            </div>
-        """)
-
-    with gr.Tabs():
-        with gr.TabItem("🔍 Diagnostic Scanner"):
-            with gr.Row():
-                with gr.Column(scale=1):
-                    input_img = gr.Image(label="Input MRI Scan", type="pil", height=400)
-                    btn = gr.Button("START NEURAL ANALYSIS", variant="primary")
-                    
-                with gr.Column(scale=1):
-                    output_html = gr.HTML(label="Diagnostic Report")
-                    overlay_img = gr.Image(label="Visualization Overlay", height=400)
-        
-        with gr.TabItem("📖 Tumor Information"):
-            gr.Markdown("""
-            ### **Predefined Tumor Classification Reference**
+        with gr.Column(scale=1):
+            img_input = gr.Image(label="Upload Patient MRI", type="pil")
+            run_btn = gr.Button("PERFORM NEURAL SCAN", variant="primary")
             
-            | Tumor Type | Description | Characteristic |
-            | :--- | :--- | :--- |
-            | **Glioma** | Tumors that start in the glial cells of the brain or spine. | Most common primary brain tumor. |
-            | **Meningioma** | A tumor that arises from the meninges (membranes covering the brain). | Usually slow-growing and benign. |
-            | **Pituitary** | Abnormal growths that develop in the pituitary gland. | Can affect hormone levels and vision. |
-            | **No Tumor** | Normal brain tissue detected. | No significant mass effect or abnormal growth seen. |
+        with gr.Column(scale=1):
+            res_html = gr.HTML(label="Diagnostic Findings", elem_id="result-text")
+            img_output = gr.Image(label="Visualization Overlay")
             
-            > **Disclaimer:** This AI is for educational and research purposes only. Always consult a certified radiologist for medical diagnosis.
-            """)
-
-    btn.click(classify_tumor, input_img, [output_html, overlay_img])
+    run_btn.click(classify_tumor, img_input, [res_html, img_output])
 
 if __name__ == "__main__":
     demo.launch()