File size: 1,003 Bytes
ce25b57 a8d5de3 ce25b57 a8d5de3 ce25b57 a8d5de3 ce25b57 a8d5de3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | import os
from ultralytics import YOLO
class ProductDetector:
def __init__(self, model_path, conf_threshold=0.25, iou_threshold=0.7, imgsz=640, device=None):
self.model = YOLO(model_path)
self.conf_threshold = conf_threshold
self.iou_threshold = iou_threshold
self.imgsz = imgsz
self.device = device if device is not None else os.getenv("YOLO_DEVICE", "cpu")
def detect(self, image_path):
results = self.model(
image_path,
conf=self.conf_threshold,
iou=self.iou_threshold,
imgsz=self.imgsz,
device=self.device,
verbose=False,
)
detections = []
for box in results[0].boxes:
x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
detections.append({
"bbox": [x1, y1, x2, y2],
"confidence": float(box.conf[0]),
"center_y": (y1 + y2) // 2
})
return detections
|