Update app.py

app.py

@@ -4,23 +4,36 @@ from ultralytics import YOLO
 import torch
 
 # ------------------------
-# تحميل الموديل (YOLOv11 classification .pt)
+# Load the YOLOv11 classification model
 # ------------------------
 MODEL_PATH = "best.pt"
 model = YOLO(MODEL_PATH)
-CLASS_NAMES = ["ax", "co", "sa"]
+
+# Mapping from model's short labels to full orientation names
+ORIENTATION_MAP = {
+    "ax": "axial",
+    "co": "coaxial",
+    "sa": "sagittal"
+}
 
 # ------------------------
-# prediction
+# Prediction function
 # ------------------------
 def predict_orientation(image: Image.Image):
-    # inference باستخدام Ultralytics
+    # Perform inference using the model
     results = model.predict(source=image, device="cpu", imgsz=224, verbose=False)
     
-    # استخراج الاحتمالات
+    # Extract probabilities
     probs = results[0].probs.data.cpu()
     pred = torch.argmax(probs)
-    orientation = CLASS_NAMES[pred.item()]
+    
+    # Get the original class label from the model
+    original_class = model.names[pred.item()]  # "ax", "co", or "sa"
+    
+    # Map to full orientation name
+    orientation = ORIENTATION_MAP.get(original_class, original_class)
+    
+    # Get confidence score
     confidence = round(probs[pred].item(), 2)
     
     return f"Orientation: {orientation} | Confidence: {confidence}"