asus
5c7c122105
This commit includes major additions and updates to the frontend and backend architectures, introducing new dataset management, live counting features, batch processing, and triage logic. Includes new UI pages, components, and API routes.
83 lines
2.6 KiB
Python
83 lines
2.6 KiB
Python
"""YOLO-based detectors for sacks and trucks.
|
|
|
|
Each detector is a single-responsibility unit (S). New model types can be
|
|
added as new classes without touching these (O).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
from ultralytics import YOLO
|
|
|
|
from src.interfaces import Detection
|
|
|
|
|
|
class SackDetector:
|
|
"""Detects sacks (and persons) using a YOLO segmentation model."""
|
|
|
|
def __init__(self, model_path: str, conf: float = 0.40) -> None:
|
|
self._model = YOLO(model_path)
|
|
self._conf = conf
|
|
|
|
def detect(self, frame: np.ndarray) -> list[Detection]:
|
|
results = self._model.predict(
|
|
frame, conf=self._conf, verbose=False
|
|
)
|
|
return self._parse(results[0])
|
|
|
|
def _parse(self, result) -> list[Detection]:
|
|
detections: list[Detection] = []
|
|
masks = result.masks
|
|
for i, box in enumerate(result.boxes):
|
|
cls_id = int(box.cls[0])
|
|
name = self._model.names[cls_id]
|
|
if name != "sack":
|
|
continue
|
|
x1, y1, x2, y2 = box.xyxy[0].tolist()
|
|
mask = None
|
|
if masks is not None and i < len(masks):
|
|
mask = masks[i].data.cpu().numpy().squeeze()
|
|
detections.append(
|
|
Detection(
|
|
bbox=(x1, y1, x2, y2),
|
|
confidence=float(box.conf[0]),
|
|
class_id=cls_id,
|
|
class_name=name,
|
|
mask=mask,
|
|
)
|
|
)
|
|
return detections
|
|
|
|
|
|
class TruckDetector:
|
|
"""Detects trucks using a YOLO detection model."""
|
|
|
|
def __init__(self, model_path_or_model: str | YOLO, conf: float = 0.50) -> None:
|
|
if isinstance(model_path_or_model, str):
|
|
self._model = YOLO(model_path_or_model)
|
|
else:
|
|
self._model = model_path_or_model
|
|
self._conf = conf
|
|
|
|
def detect(self, frame: np.ndarray) -> list[Detection]:
|
|
results = self._model.predict(
|
|
frame, conf=self._conf, verbose=False
|
|
)
|
|
return self._parse(results[0])
|
|
|
|
def _parse(self, result) -> list[Detection]:
|
|
detections: list[Detection] = []
|
|
for box in result.boxes:
|
|
cls_id = int(box.cls[0])
|
|
name = self._model.names[cls_id]
|
|
x1, y1, x2, y2 = box.xyxy[0].tolist()
|
|
detections.append(
|
|
Detection(
|
|
bbox=(x1, y1, x2, y2),
|
|
confidence=float(box.conf[0]),
|
|
class_id=cls_id,
|
|
class_name=name,
|
|
)
|
|
)
|
|
return detections
|