Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import sahi
 import torch
-from ultralyticsplus import YOLO, render_model_output
+from ultralytics import YOLO
 
 # Download sample images
 sahi.utils.file.download_from_url(
@@ -60,26 +60,16 @@ def yolov8_inference(
     model.overrides["iou"] = iou_threshold
     
     # Perform model prediction
-    results = model.predict(image, imgsz=image_size, return_outputs=True)
+    results = model(image)
 
     # Initialize an empty list to store the output
     output = []
     
     # Iterate over the results
-    for result in results:
-        # Check if segmentation masks are available
-        if 'masks' in result and result['masks'] is not None:
-            masks = result['masks']['data']
-            for i, (mask, box) in enumerate(zip(masks, result['boxes'])):
-                label = model.names[int(result['boxes']['cls'][i])]
-                mask_coords = mask.tolist()  # Convert mask coordinates to list format
-                output.append({"label": label, "mask_coords": mask_coords})
-        else:
-            # If masks are not available, just extract bounding box information
-            for i, box in enumerate(result['boxes']):
-                label = model.names[int(result['boxes']['cls'][i])]
-                bbox = box['xyxy'].tolist()  # Bounding box coordinates
-                output.append({"label": label, "bbox_coords": bbox})
+    for i,box in enumerate(results[0].boxes):
+        label = results[0].names[box.cls[0].item()]
+        bbox = box.xyxy[0]
+        output.append({"label": label, "bbox_coords": bbox})
     
     return output