Update app.py

app.py

@@ -1,37 +1,24 @@
 # app.py for Hugging Face Spaces
 
 import gradio as gr
-from ultralytics import YOLO
-import torch
-import ultralytics.nn.tasks
-from huggingface_hub import hf_hub_download
+from ultralyticsplus import YOLO
 from PIL import Image
 import numpy as np
 
-# Patch torch_safe_load to use weights_only=False
-def patched_torch_safe_load(weight):
-    try:
-        return torch.load(weight, map_location='cpu', weights_only=False)
-    except Exception as e:
-        raise e
-
-ultralytics.nn.tasks.torch_safe_load = patched_torch_safe_load
-
-# Download and load the model
-model_path = hf_hub_download(repo_id="foduucom/product-detection-in-shelf-yolov8", filename="best.pt")
-model = YOLO(model_path)
+# Load the model
+model = YOLO('foduucom/product-detection-in-shelf-yolov8')
 
 # Set model parameters
-model.conf = 0.25  # confidence threshold (note: in newer ultralytics, it's in predict)
-model.iou = 0.45  # IoU threshold
-model.agnostic_nms = False
-model.max_det = 1000
+model.overrides['conf'] = 0.25  # NMS confidence threshold
+model.overrides['iou'] = 0.45  # NMS IoU threshold
+model.overrides['agnostic_nms'] = False  # NMS class-agnostic
+model.overrides['max_det'] = 1000  # maximum number of detections per image
 
 def detect_skus(image):
-    # image is PIL Image from Gradio
-    results = model(image)  # YOLO accepts PIL
+    # Convert Gradio image to numpy array if needed, but YOLO accepts PIL
+    results = model(image)
     
-    # Extract unique SKU names
+    # Extract unique SKU names (assuming classes are SKU names)
     sku_set = set()
     for result in results:
         for box in result.boxes: