feat: add counting bench, triage, and dataset modules
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.
This commit is contained in:
1 parent
8285400254
commit
5c7c122105
80 files changed
+20074
-1412
No files matched your search
@@ -0,0 +1,82 @@
|
||||
"""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
|
||||
Reference in new issue
Block a user