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.
468 lines
19 KiB
Python
468 lines
19 KiB
Python
"""Live counting test bench: point a trained model at an RTSP stream and watch it count.
|
|
|
|
This is a **test harness**, not the production counter. It reuses the real
|
|
pipeline pieces from `algoritma-batch` — ByteTrack, the bbox stabiliser and the
|
|
line-cross counter with its spatial dedup — so what you see here is what
|
|
`predict.py` would do. What it deliberately leaves out is everything stateful:
|
|
no batch lifecycle, no SQLite, no truck-presence state machine. The question it
|
|
answers is "does this model count correctly on this camera", and those parts
|
|
only get in the way of answering it.
|
|
|
|
One session at a time, holding the GPU lock, because the GPU is shared with
|
|
training and auto-annotation (REQ-065, REQ-070).
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import threading
|
|
import time
|
|
from typing import Optional
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from backend import config, jobs
|
|
|
|
# `algoritma-batch/src` is copied to /app/src in the image; in a source checkout
|
|
# it still lives under algoritma-batch/. Both are made importable as `src.*`.
|
|
_REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
for candidate in ("/app", os.path.join(_REPO, "algoritma-batch")):
|
|
if os.path.isdir(os.path.join(candidate, "src")) and candidate not in sys.path:
|
|
sys.path.insert(0, candidate)
|
|
|
|
|
|
class LiveCountError(Exception):
|
|
pass
|
|
|
|
|
|
class Session:
|
|
"""One running counter. Owns a capture thread and the latest rendered frame."""
|
|
|
|
def __init__(self, source: str, model_path: str, line_y: int,
|
|
line_x_start: int, line_x_end: int, conf: float,
|
|
dedup_radius: float, margin: int, imgsz: int,
|
|
entry_travel_min: float, handoff_radius: float,
|
|
unload_confirm_frames: int, min_area_scale: float,
|
|
spatial_dedup: bool):
|
|
self.source = source
|
|
self.model_path = model_path
|
|
self.line_y = line_y
|
|
self.line_x_start = line_x_start
|
|
self.line_x_end = line_x_end
|
|
self.conf = conf
|
|
self.dedup_radius = dedup_radius
|
|
self.margin = margin
|
|
self.imgsz = imgsz
|
|
self.entry_travel_min = entry_travel_min
|
|
self.handoff_radius = handoff_radius
|
|
self.unload_confirm_frames = unload_confirm_frames
|
|
self.min_area_scale = min_area_scale
|
|
self.spatial_dedup = spatial_dedup
|
|
|
|
self.started_at = time.time()
|
|
self.error = ""
|
|
self.stopping = False
|
|
self.frames = 0
|
|
self.fps = 0.0
|
|
self.loading = 0
|
|
self.unloading = 0
|
|
self.tracked = 0
|
|
self.ignored = 0
|
|
self.too_small = 0
|
|
self.traced = 0
|
|
self.trace_path = os.path.join(
|
|
config.DATA_DIR, "live-count", f"session-{int(self.started_at)}.jsonl")
|
|
self.events: list[dict] = []
|
|
|
|
self._jpeg: Optional[bytes] = None
|
|
self._counter = None # set once the worker builds it
|
|
self._lock = threading.Lock()
|
|
self._thread = threading.Thread(target=self._run, name="live-count", daemon=True)
|
|
|
|
# -- public -----------------------------------------------------------
|
|
|
|
def start(self) -> None:
|
|
self._thread.start()
|
|
|
|
def stop(self) -> None:
|
|
self.stopping = True
|
|
self._thread.join(timeout=15)
|
|
|
|
def snapshot(self) -> Optional[bytes]:
|
|
with self._lock:
|
|
return self._jpeg
|
|
|
|
def move_line(self, line_y: Optional[int] = None, line_x_start: Optional[int] = None,
|
|
line_x_end: Optional[int] = None) -> dict:
|
|
"""Reposition the counting line without restarting.
|
|
|
|
Placing a line correctly means watching the stream while you move it,
|
|
which is impossible if moving it requires a restart — and a restart
|
|
throws the counts away. Already-counted tracks keep their verdict; the
|
|
counter only decides for a track the first time it crosses.
|
|
"""
|
|
if line_y is not None:
|
|
self.line_y = max(0, min(720, int(line_y)))
|
|
if line_x_start is not None:
|
|
self.line_x_start = max(0, min(1280, int(line_x_start)))
|
|
if line_x_end is not None:
|
|
self.line_x_end = max(0, min(1280, int(line_x_end)))
|
|
if self.line_x_end < self.line_x_start:
|
|
self.line_x_start, self.line_x_end = self.line_x_end, self.line_x_start
|
|
if self._counter is not None:
|
|
self._counter.line_y = self.line_y
|
|
self._counter.line_x_start = self.line_x_start
|
|
self._counter.line_x_end = self.line_x_end
|
|
return {"y": self.line_y, "x_start": self.line_x_start, "x_end": self.line_x_end}
|
|
|
|
def status(self) -> dict:
|
|
return {
|
|
"running": self._thread.is_alive(),
|
|
"source": self.source,
|
|
"model_path": self.model_path,
|
|
"error": self.error,
|
|
"frames": self.frames,
|
|
"fps": round(self.fps, 1),
|
|
"loading": self.loading,
|
|
"unloading": self.unloading,
|
|
"net": self.loading - self.unloading,
|
|
"tracked": self.tracked,
|
|
"ignored": self.ignored,
|
|
"elapsed": round(time.time() - self.started_at, 1),
|
|
"line": {"y": self.line_y, "x_start": self.line_x_start, "x_end": self.line_x_end},
|
|
"too_small": self.too_small,
|
|
"traced": self.traced,
|
|
"trace_path": self.trace_path,
|
|
"events": self.events[-25:],
|
|
}
|
|
|
|
# -- worker -----------------------------------------------------------
|
|
|
|
def _run(self) -> None:
|
|
# The GPU is shared. Waiting here rather than failing means "start" is
|
|
# safe to press while a training run is finishing.
|
|
if not jobs.gpu_lock.acquire(timeout=30):
|
|
busy = jobs.running_types()
|
|
self.error = f"GPU busy with a {busy[0] if busy else 'background'} job"
|
|
return
|
|
capture = None
|
|
try:
|
|
from ultralytics import YOLO
|
|
from src.counting import LineCrossCounter
|
|
from src.stabilizer import BboxStabilizer
|
|
from src.tracking import ByteTrackTracker
|
|
|
|
model = YOLO(self.model_path)
|
|
# Warm-up: the first CUDA call inside the tracker has been seen to
|
|
# segfault without it (same reason predict.py does this).
|
|
model(np.zeros((720, 1280, 3), dtype=np.uint8), imgsz=self.imgsz, verbose=False)
|
|
|
|
tracker = ByteTrackTracker(model, self.conf)
|
|
stabilizer = BboxStabilizer(ema_alpha=0.35, max_hold_frames=10,
|
|
max_height_ratio=1.5, min_height_ratio=0.70)
|
|
counter = LineCrossCounter(
|
|
line_y=self.line_y, line_x_start=self.line_x_start,
|
|
line_x_end=self.line_x_end, margin=self.margin,
|
|
dedup_radius=self.dedup_radius,
|
|
entry_travel_min=self.entry_travel_min,
|
|
handoff_radius=self.handoff_radius,
|
|
unload_confirm_frames=self.unload_confirm_frames,
|
|
spatial_dedup=self.spatial_dedup,
|
|
)
|
|
self._counter = counter
|
|
|
|
capture = _open(self.source)
|
|
if capture is None or not capture.isOpened():
|
|
raise LiveCountError(f"Could not open source: {self.source}")
|
|
|
|
tick = time.time()
|
|
since = 0
|
|
while not self.stopping:
|
|
ok, frame = capture.read()
|
|
if not ok or frame is None:
|
|
# A file simply ends. On a stream this only means the
|
|
# decoder has not produced a new frame yet, so wait briefly
|
|
# — long enough not to spin, short enough not to become the
|
|
# new frame-rate ceiling.
|
|
if _is_stream(self.source):
|
|
time.sleep(0.005)
|
|
continue
|
|
break
|
|
|
|
frame = cv2.resize(frame, (1280, 720))
|
|
detections = [d for d in tracker.update(frame, []) if d.class_name == "sack"]
|
|
stable = stabilizer.update(detections)
|
|
inside, outside = [], []
|
|
small = 0
|
|
for det in stable:
|
|
x1, y1, x2, y2 = det.bbox
|
|
centre_x = (x1 + x2) / 2
|
|
# Perspective-aware area gate, the curve `predict.py` uses:
|
|
# a box that small at that depth is a fragment, not a sack.
|
|
if _too_small(det.bbox, self.min_area_scale):
|
|
small += 1
|
|
outside.append(det)
|
|
elif self.line_x_start <= centre_x <= self.line_x_end:
|
|
inside.append(det)
|
|
else:
|
|
outside.append(det)
|
|
self.too_small = small
|
|
for event in counter.update(inside):
|
|
self.events.append({
|
|
"track_id": event.get("track_id"),
|
|
"direction": event.get("direction", "loading"),
|
|
"at": round(time.time() - self.started_at, 1),
|
|
})
|
|
|
|
self._write_traces(counter.drain_traces())
|
|
self.loading = counter.loading_count
|
|
self.unloading = counter.unloading_count
|
|
self.tracked = len(inside)
|
|
self.ignored = len(outside)
|
|
self.frames += 1
|
|
since += 1
|
|
|
|
now = time.time()
|
|
if now - tick >= 1.0:
|
|
self.fps = since / (now - tick)
|
|
tick, since = now, 0
|
|
|
|
self._render(frame, inside, outside, counter)
|
|
except Exception as exc: # surfaced in status(), not swallowed
|
|
self.error = f"{type(exc).__name__}: {exc}"
|
|
finally:
|
|
# The lock is released no matter what tearing down the capture does.
|
|
# It was the other order once, and one exception in release() leaked
|
|
# the GPU for the lifetime of the process.
|
|
try:
|
|
if capture is not None:
|
|
capture.release()
|
|
except Exception as exc:
|
|
if not self.error:
|
|
self.error = f"capture release failed: {exc}"
|
|
finally:
|
|
jobs.gpu_lock.release()
|
|
|
|
def _write_traces(self, traces: list) -> None:
|
|
"""Append finished tracks to a JSONL, one object per track.
|
|
|
|
This is the file that answers "was that the model, the tracker or the
|
|
counter" on a clip with a known count: every track that ever existed
|
|
lands here with its trajectory and the reason it did or did not count.
|
|
"""
|
|
if not traces:
|
|
return
|
|
import json
|
|
try:
|
|
os.makedirs(os.path.dirname(self.trace_path), exist_ok=True)
|
|
with open(self.trace_path, "a", encoding="utf-8") as handle:
|
|
for record in traces:
|
|
handle.write(json.dumps(record, default=str) + "\n")
|
|
self.traced += len(traces)
|
|
except OSError as exc:
|
|
if not self.error:
|
|
self.error = f"trace write failed: {exc}"
|
|
|
|
def _render(self, frame, detections, ignored, counter) -> None:
|
|
height, width = frame.shape[:2]
|
|
|
|
# Shade what the region excludes. Without this the neighbouring truck's
|
|
# sacks simply vanish from the overlay, and "are they being ignored?"
|
|
# looks identical to "is the model missing them?".
|
|
if self.line_x_start > 0 or self.line_x_end < width:
|
|
shade = frame.copy()
|
|
if self.line_x_start > 0:
|
|
cv2.rectangle(shade, (0, 0), (self.line_x_start, height), (0, 0, 0), -1)
|
|
if self.line_x_end < width:
|
|
cv2.rectangle(shade, (self.line_x_end, 0), (width, height), (0, 0, 0), -1)
|
|
cv2.addWeighted(shade, 0.55, frame, 0.45, 0, frame)
|
|
|
|
# Ignored detections stay visible, in grey, so the region can be judged.
|
|
for det in ignored:
|
|
x1, y1, x2, y2 = (int(v) for v in det.bbox)
|
|
cv2.rectangle(frame, (x1, y1), (x2, y2), (130, 130, 130), 1)
|
|
|
|
for edge in (self.line_x_start, self.line_x_end):
|
|
if 0 < edge < width:
|
|
cv2.line(frame, (edge, 0), (edge, height), (255, 0, 255), 2)
|
|
cv2.putText(frame, "IGNORED", (max(4, self.line_x_start - 92), height - 14),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 255), 1, cv2.LINE_AA)
|
|
cv2.putText(frame, "IGNORED", (min(width - 88, self.line_x_end + 8), height - 14),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 255), 1, cv2.LINE_AA)
|
|
|
|
for det in detections:
|
|
x1, y1, x2, y2 = (int(v) for v in det.bbox)
|
|
counted = bool(counter.counted_tracks.get(det.track_id))
|
|
colour = (74, 222, 128) if counted else (248, 191, 113)
|
|
cv2.rectangle(frame, (x1, y1), (x2, y2), colour, 2)
|
|
cv2.putText(frame, f"#{det.track_id} {det.confidence:.2f}", (x1, max(14, y1 - 6)),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.45, colour, 1, cv2.LINE_AA)
|
|
|
|
cv2.line(frame, (self.line_x_start, self.line_y), (self.line_x_end, self.line_y),
|
|
(0, 255, 255), 2)
|
|
for edge in (self.line_y - self.margin, self.line_y + self.margin):
|
|
cv2.line(frame, (self.line_x_start, edge), (self.line_x_end, edge),
|
|
(0, 160, 160), 1)
|
|
|
|
panel = f"IN {self.loading} OUT {self.unloading} NET {self.loading - self.unloading}"
|
|
cv2.rectangle(frame, (12, 12), (12 + 9 * len(panel) + 20, 84), (0, 0, 0), -1)
|
|
cv2.putText(frame, panel, (24, 46), cv2.FONT_HERSHEY_SIMPLEX, 0.8,
|
|
(74, 222, 128), 2, cv2.LINE_AA)
|
|
cv2.putText(frame, f"{self.fps:.1f} fps {self.tracked} tracked {self.ignored} ignored",
|
|
(24, 72), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (200, 200, 200), 1, cv2.LINE_AA)
|
|
|
|
ok, buffer = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 75])
|
|
if ok:
|
|
with self._lock:
|
|
self._jpeg = buffer.tobytes()
|
|
|
|
|
|
def _too_small(bbox, scale: float) -> bool:
|
|
"""`predict.py`'s perspective curve: a sack at the top of the frame is
|
|
genuinely smaller in pixels than the same sack at the bottom, so one flat
|
|
threshold either lets fragments through up close or discards real sacks far
|
|
away. Interpolated in the 1280x720 space the pipeline works in. `scale` of 0
|
|
turns the gate off."""
|
|
if scale <= 0:
|
|
return False
|
|
x1, y1, x2, y2 = bbox
|
|
centre_y = (y1 + y2) / 2.0
|
|
top_y, bottom_y = 133.0, 720.0
|
|
top_area, bottom_area = 3556.0, 11111.0
|
|
if centre_y <= top_y:
|
|
minimum = top_area
|
|
elif centre_y >= bottom_y:
|
|
minimum = bottom_area
|
|
else:
|
|
ratio = (centre_y - top_y) / (bottom_y - top_y)
|
|
minimum = top_area + ratio * (bottom_area - top_area)
|
|
return (x2 - x1) * (y2 - y1) < minimum * scale
|
|
|
|
|
|
def _is_stream(source: str) -> bool:
|
|
return str(source).startswith(("rtsp://", "rtmp://", "http://", "https://"))
|
|
|
|
|
|
class _ThreadedStream:
|
|
"""Decode in a background thread and always hand out the newest frame.
|
|
|
|
A plain VideoCapture.read() on RTSP is blocking, and decoding 1080p costs
|
|
more than inference does — measured here at 6.7 fps end-to-end against 200
|
|
fps for the model itself. Worse, reading slower than the camera sends builds
|
|
a backlog, so the picture drifts further behind real time the longer it
|
|
runs. Dropping stale frames keeps latency flat, which is what a counting
|
|
test needs to mean anything. `predict.py` does the same thing.
|
|
"""
|
|
|
|
def __init__(self, source: str):
|
|
self._capture = cv2.VideoCapture(source)
|
|
try:
|
|
self._capture.set(cv2.CAP_PROP_BUFFERSIZE, 1)
|
|
except Exception:
|
|
pass
|
|
self._frame = None
|
|
self._lock = threading.Lock()
|
|
self._running = True
|
|
self._thread = threading.Thread(target=self._pump, daemon=True)
|
|
if self._capture.isOpened():
|
|
self._thread.start()
|
|
|
|
def _pump(self) -> None:
|
|
while self._running:
|
|
ok, frame = self._capture.read()
|
|
if not ok:
|
|
time.sleep(0.01)
|
|
continue
|
|
with self._lock:
|
|
self._frame = frame
|
|
|
|
def isOpened(self) -> bool:
|
|
return self._capture.isOpened()
|
|
|
|
def read(self):
|
|
with self._lock:
|
|
if self._frame is None:
|
|
return False, None
|
|
frame, self._frame = self._frame, None
|
|
return True, frame
|
|
|
|
def release(self) -> None:
|
|
self._running = False
|
|
# The thread is only started when the capture opened, so a failed
|
|
# source would otherwise raise "cannot join thread before it is
|
|
# started" here — inside the caller's finally, skipping the GPU lock
|
|
# release and wedging every later session on "GPU busy".
|
|
if self._thread.is_alive():
|
|
self._thread.join(timeout=2)
|
|
self._capture.release()
|
|
|
|
|
|
def _open(source: str):
|
|
if _is_stream(source):
|
|
os.environ.setdefault(
|
|
"OPENCV_FFMPEG_CAPTURE_OPTIONS",
|
|
"rtsp_transport;tcp|buffer_size;20480000|max_delay;500000",
|
|
)
|
|
return _ThreadedStream(source)
|
|
return cv2.VideoCapture(source)
|
|
|
|
|
|
# ---- module-level single session ----------------------------------------
|
|
|
|
_session: Optional[Session] = None
|
|
_guard = threading.Lock()
|
|
|
|
|
|
def start(source: str, model_path: str, line_y: int, line_x_start: int, line_x_end: int,
|
|
conf: float = 0.35, dedup_radius: float = 60.0, margin: int = 5,
|
|
imgsz: int = 640, entry_travel_min: float = 60.0,
|
|
handoff_radius: float = 100.0, unload_confirm_frames: int = 3,
|
|
min_area_scale: float = 1.0, spatial_dedup: bool = False) -> dict:
|
|
global _session
|
|
with _guard:
|
|
if _session is not None and _session.status()["running"]:
|
|
raise LiveCountError("A counting session is already running — stop it first")
|
|
if not os.path.isfile(model_path):
|
|
raise LiveCountError(f"Model not found: {model_path}")
|
|
if not _is_stream(source) and not os.path.isfile(source):
|
|
raise LiveCountError(f"Source not found: {source}")
|
|
_session = Session(source, model_path, line_y, line_x_start, line_x_end,
|
|
conf, dedup_radius, margin, imgsz, entry_travel_min,
|
|
handoff_radius, unload_confirm_frames, min_area_scale,
|
|
spatial_dedup)
|
|
_session.start()
|
|
time.sleep(0.4) # let an immediate failure surface in the response
|
|
return _session.status()
|
|
|
|
|
|
def stop() -> dict:
|
|
global _session
|
|
with _guard:
|
|
if _session is None:
|
|
return {"running": False}
|
|
_session.stop()
|
|
report = _session.status()
|
|
_session = None
|
|
return report
|
|
|
|
|
|
def move_line(line_y=None, line_x_start=None, line_x_end=None) -> dict:
|
|
if _session is None:
|
|
raise LiveCountError("No counting session is running")
|
|
return _session.move_line(line_y, line_x_start, line_x_end)
|
|
|
|
|
|
def status() -> dict:
|
|
if _session is None:
|
|
# Same shape as a live session, so callers never branch on presence.
|
|
return {"running": False, "loading": 0, "unloading": 0, "net": 0, "frames": 0,
|
|
"fps": 0.0, "tracked": 0, "ignored": 0, "too_small": 0, "traced": 0,
|
|
"trace_path": "", "elapsed": 0.0, "events": [],
|
|
"error": "", "source": "", "model_path": "",
|
|
"line": {"y": 0, "x_start": 0, "x_end": 1280}}
|
|
return _session.status()
|
|
|
|
|
|
def snapshot() -> Optional[bytes]:
|
|
return _session.snapshot() if _session is not None else None
|