Update app.py

app.py

@@ -1,46 +1,34 @@
 import gradio as gr
 from PIL import Image
+from ultralytics import YOLO
 import torch
-from pathlib import Path
 
 # ------------------------
-# تحميل الموديل
-# ------------------------
-MODEL_PATH = Path("best.pt")
-model = torch.load(MODEL_PATH, map_location="cpu",weights_only=False)
-model.eval()
-
+MODEL_PATH = "best.pt"
+model = YOLO(MODEL_PATH)  
 CLASS_NAMES = ["ax", "co", "sa"]
 
 # ------------------------
 #  prediction
 # ------------------------
 def predict_orientation(image: Image.Image):
-    import torchvision.transforms as transforms
-    transform = transforms.Compose([
-        transforms.Resize((224,224)),  
-        transforms.ToTensor(),
-    ])
-    img_tensor = transform(image).unsqueeze(0)  # batch dimension
-
-    with torch.no_grad():
-        outputs = model(img_tensor)
-        probs = torch.softmax(outputs, dim=1)
-        conf, pred = torch.max(probs, 1)
-        orientation = CLASS_NAMES[pred.item()]
-        confidence = round(conf.item(), 2)
+    results = model.predict(source=image, imgsz=224, device="cpu")  
+    probs = results[0].probs  
+    pred = torch.argmax(probs)
+    orientation = CLASS_NAMES[pred.item()]
+    confidence = round(probs[pred].item(), 2)
     
     return f"Orientation: {orientation} | Confidence: {confidence}"
 
 # ------------------------
-#  Gradio
+#  Gradio Interface
 # ------------------------
 iface = gr.Interface(
     fn=predict_orientation,
     inputs=gr.Image(type="pil"),
     outputs="text",
     title="MRI Orientation Predictor",
-    description="upload your image and the model output prediction and confidence"
+    description="Upload your image and the model outputs prediction and confidence."
 )
 
 iface.launch()