Update app.py

app.py

@@ -1,36 +1,30 @@
 import gradio as gr
 from huggingface_hub import hf_hub_download
-import torch
+from ultralytics import YOLO
 from PIL import Image
 
-# Load your model from your HF repo
-model = torch.load(
-    hf_hub_download(repo_id="newtechdevng/detect", filename="best.pt"),
-    map_location="cpu"
-)
-model.eval()
-
-# Define labels
-labels = ["car", "bike", "mountain", "road"]  # update to match your model
+# Load YOLO model correctly
+model_path = hf_hub_download(repo_id="newtechdevng/detect", filename="best.pt")
+model = YOLO(model_path)
 
 def predict(image):
-    # Preprocess
-    from torchvision import transforms
-    transform = transforms.Compose([
-        transforms.Resize((224, 224)),
-        transforms.ToTensor(),
-    ])
-    tensor = transform(image).unsqueeze(0)
+    results = model(image)
+    result = results[0]
+    
+    output = []
+    for box in result.boxes:
+        label = result.names[int(box.cls)]
+        confidence = float(box.conf)
+        output.append(f"{label}: {confidence:.2f}")
     
-    with torch.no_grad():
-        output = model(tensor)
-        predicted = torch.argmax(output, dim=1).item()
+    if not output:
+        return "No objects detected"
     
-    return labels[predicted]
+    return "\n".join(output)
 
 gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
-    outputs=gr.Label(),
-    title="Car/Bike/Mountain/Road Detector"
+    outputs=gr.Text(label="Detections"),
+    title="Car / Bike / Mountain / Road Detector"
 ).launch()