Update app.py

app.py

@@ -9,23 +9,71 @@ model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
 model.eval()
 
 def classify_mri(image):
+    if image is None:
+        return None, "Please upload an MRI image."
+
     inputs = processor(images=image, return_tensors="pt")
 
     with torch.no_grad():
         outputs = model(**inputs)
         probabilities = torch.softmax(outputs.logits, dim=-1)[0]
 
-    return {
+    scores = {
         model.config.id2label[i]: float(probabilities[i])
         for i in range(len(probabilities))
     }
 
-demo = gr.Interface(
-    fn=classify_mri,
-    inputs=gr.Image(type="pil", label="Upload brain MRI"),
-    outputs=gr.Label(num_top_classes=4),
-    title="Brain MRI Tumour Classification",
-    description="Educational prototype only. Not for medical diagnosis."
-)
+    prediction = max(scores, key=scores.get)
+    confidence = scores[prediction]
+
+    result = (
+        f"### Prediction: {prediction}\n\n"
+        f"Confidence: **{confidence:.1%}**\n\n"
+        "This is an educational prototype and not a medical diagnosis."
+    )
+
+    return scores, result
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # Brain MRI Tumour Classification
+
+        Upload a brain MRI image to classify it as:
+
+        **Glioma · Meningioma · No Tumour · Pituitary**
+        """
+    )
+
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(
+                type="pil",
+                label="Upload MRI Image"
+            )
+            classify_button = gr.Button(
+                "Classify MRI",
+                variant="primary"
+            )
+
+        with gr.Column():
+            prediction_output = gr.Label(
+                label="Class Probabilities",
+                num_top_classes=4
+            )
+            result_output = gr.Markdown()
+
+    classify_button.click(
+        fn=classify_mri,
+        inputs=image_input,
+        outputs=[prediction_output, result_output]
+    )
+
+    gr.Markdown(
+        """
+        **Disclaimer:** This application is for educational purposes only.
+        It must not be used for clinical diagnosis or treatment decisions.
+        """
+    )
 
 demo.launch()