Spaces:
Sleeping
Sleeping
File size: 1,500 Bytes
da1e434 |
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 39 40 41 42 43 44 45 46 47 48 |
"""YOLOv8 wrapper for ball, stump, and pad detection."""
from ultralytics import YOLO
from box import Box
from utils.io_utils import logger
import yaml
from pathlib import Path
# -------------------------------------------------
# Robust path to lbw_drs_ai/config/config.yaml
# -------------------------------------------------
ROOT = Path(__file__).resolve().parents[1] # …/lbw_drs_ai
CFG_PATH = ROOT / "config" / "config.yaml"
if not CFG_PATH.is_file():
raise FileNotFoundError(f"Config file not found: {CFG_PATH}")
CONFIG = Box(yaml.safe_load(CFG_PATH.read_text()))
# -------------------------------------------------
CLASS_MAP = {
0: "ball",
1: "stump",
2: "pad",
}
class Detector:
def __init__(self):
logger.info("Loading YOLOv8 model…")
self.model = YOLO(CONFIG.model.yolo_weights)
self.conf = CONFIG.model.conf_threshold
self.iou = CONFIG.model.iou_threshold
def infer(self, frame):
"""Return list of detections as dicts: {cls, conf, bbox}."""
results = self.model.predict(frame, conf=self.conf, iou=self.iou, verbose=False)
detections = []
for r in results:
for b in r.boxes:
cls_id = int(b.cls)
detections.append({
"class": CLASS_MAP.get(cls_id, str(cls_id)),
"conf": float(b.conf),
"bbox": b.xyxy[0].cpu().numpy(), # [x1,y1,x2,y2]
})
return detections
|