Update app.py

app.py

@@ -1,18 +1,31 @@
 import gradio as gr
+import spaces
 import torch
 from transformers import AutoImageProcessor, AutoModelForImageClassification
 
 MODEL_ID = "DiggaDee/brain-tumor-resnet18"
 
 processor = AutoImageProcessor.from_pretrained(MODEL_ID)
+
 model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
+model.to("cuda")
 model.eval()
 
+
+@spaces.GPU
 def classify_mri(image):
     if image is None:
-        return None, "Please upload an MRI image."
+        return {}, "Please upload an MRI image."
 
-    inputs = processor(images=image, return_tensors="pt")
+    inputs = processor(
+        images=image,
+        return_tensors="pt"
+    )
+
+    inputs = {
+        key: value.to("cuda")
+        for key, value in inputs.items()
+    }
 
     with torch.no_grad():
         outputs = model(**inputs)
@@ -23,17 +36,19 @@ def classify_mri(image):
         for i in range(len(probabilities))
     }
 
-    prediction = max(scores, key=scores.get)
-    confidence = scores[prediction]
+    prediction_index = probabilities.argmax().item()
+    prediction = model.config.id2label[prediction_index]
+    confidence = probabilities[prediction_index].item()
 
     result = (
         f"### Prediction: {prediction}\n\n"
         f"Confidence: **{confidence:.1%}**\n\n"
-        "This is an educational prototype and not a medical diagnosis."
+        "Educational prototype only. Not for medical diagnosis."
     )
 
     return scores, result
 
+
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
@@ -51,6 +66,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 type="pil",
                 label="Upload MRI Image"
             )
+
             classify_button = gr.Button(
                 "Classify MRI",
                 variant="primary"
@@ -61,19 +77,24 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 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]
+        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.
+        **Disclaimer:** This application is for educational purposes only
+        and must not be used for clinical diagnosis.
         """
     )
 
+
 demo.launch()