Update app.py

app.py

@@ -1,32 +1,33 @@
 import gradio as gr
 from PIL import Image
-from ultralytics import YOLO
+from ultralytics.nn.tasks import ClassificationModel
 import torch
 
-# ------------------------
-# تحميل الموديل التصنيفي
-# ------------------------
 MODEL_PATH = "best.pt"
-model = YOLO(MODEL_PATH)  # موديل تصنيف من Ultralytics
+model = ClassificationModel(MODEL_PATH)
+model.eval()
+
 CLASS_NAMES = ["ax", "co", "sa"]
 
 # ------------------------
 # prediction
 # ------------------------
 def predict_orientation(image: Image.Image):
-    # نستخدم predict على صورة PIL
-    results = model.predict(source=image, imgsz=224, device="cpu", verbose=False)
-    
-    # Ultralytics ClassificationModel بيرجع probs لكل class
-    # التأكد من وجود probs
-    if hasattr(results[0], "probs"):
-        probs = results[0].probs  # tensor
-        pred = torch.argmax(probs)
+    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(probs[pred].item(), 2)
-        return f"Orientation: {orientation} | Confidence: {confidence}"
-    else:
-        return "Error: Could not get probabilities from model."
+        confidence = round(conf.item(), 2)
+
+    return f"Orientation: {orientation} | Confidence: {confidence}"
 
 # ------------------------
 # Gradio Interface