Update app.py

app.py

@@ -1,32 +1,28 @@
 import gradio as gr
 from PIL import Image
-from ultralytics.nn.tasks import ClassificationModel
+from ultralytics import YOLO
 import torch
 
+# ------------------------
+# تحميل الموديل (YOLOv11 classification .pt)
+# ------------------------
 MODEL_PATH = "best.pt"
-model = ClassificationModel(MODEL_PATH)
-model.eval()
-
+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)
-
+    # inference باستخدام Ultralytics
+    results = model.predict(source=image, device="cpu", imgsz=224, verbose=False)
+    
+    # استخراج الاحتمالات
+    probs = results[0].probs.data.cpu()
+    pred = torch.argmax(probs)
+    orientation = CLASS_NAMES[pred.item()]
+    confidence = round(probs[pred].item(), 2)
+    
     return f"Orientation: {orientation} | Confidence: {confidence}"
 
 # ------------------------