src/detection.py

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