Update app.py

app.py

@@ -1,16 +1,32 @@
 import gradio as gr
+from ultralytics import YOLO
 from PIL import Image, ImageDraw
 import io
 from fpdf import FPDF
 import tempfile
 import cv2
 import os
-import torch
+import torch.serialization
 import numpy as np
+from urllib.request import urlretrieve
+from ultralytics.nn.modules import Conv, Bottleneck, SPPF, C2f, Detect
+from ultralytics.nn.tasks import DetectionModel
+from torch.nn.modules.container import Sequential
 
-# Manually load the model with weights_only=False to avoid UnpicklingError
-ckpt = torch.load('model/yolov8n.pt', map_location='cpu', weights_only=False)
-model = ckpt['model']
+# Allowlist all necessary YOLOv8 modules to avoid UnpicklingError
+torch.serialization.add_safe_globals([DetectionModel, Sequential, Conv, Bottleneck, SPPF, C2f, Detect])
+
+# Model path
+model_path = 'model/yolov8n.pt'
+
+# Download model if not found
+if not os.path.exists(model_path):
+    os.makedirs('model', exist_ok=True)
+    url = 'https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt'
+    urlretrieve(url, model_path)
+
+# Load the YOLOv8 model
+model = YOLO(model_path)
 
 def detect_empty_spots(img_width, img_height, bottle_bboxes, min_gap=50):
     empty_bboxes = []
@@ -40,8 +56,8 @@ def process_input(input_file):
     else:
         process_img_path = input_file
 
-    # Run inference with parameters
-    results = model(process_img_path, conf=0.25, iou=0.45, agnostic_nms=False, max_det=1000)
+    # Run inference with YOLO parameters
+    results = model.predict(process_img_path, conf=0.25, iou=0.45, agnostic_nms=False, max_det=1000)
     boxes = results[0].boxes
     class_names = results[0].names