andrew/reTraining
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: sync live-count, exemplar annotation modules, and update .gitignore

Browse Source
This commit is contained in:
ervanfahriaw committed 2026-08-24 14:57:42 +07:00
1 parent b6624eeff9
commit ac95674c07
39 files changed
+3979 -444
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

backend/live_count.py
+54 -128
View File
@@ -1,4 +1,4 @@
"""Live counting test bench: point a trained model at an RTSP stream and watch it count.
"""Live counting test bench: point a trained model at a live 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
@@ -22,6 +22,10 @@ import cv2
import numpy as np
from backend import config, jobs
from backend import live_render
from backend.live_source import ( # re-exported: callers catch live_count.LiveCountError
LiveCountError, _is_stream, _open, whep_to_rtsp, whep_url,
)
# `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.*`.
@@ -31,10 +35,6 @@ for candidate in ("/app", os.path.join(_REPO, "algoritma-batch")):
sys.path.insert(0, candidate)
class LiveCountError(Exception):
pass
class Session:
"""One running counter. Owns a capture thread and the latest rendered frame."""
@@ -43,8 +43,11 @@ class Session:
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):
spatial_dedup: bool, whep_url: str = ""):
self.source = source
# When the browser watches the camera over WebRTC it never asks for the
# MJPEG, so encoding a JPEG per frame would be pure waste (REQ-177).
self.whep_url = whep_url
self.model_path = model_path
self.line_y = line_y
self.line_x_start = line_x_start
@@ -75,6 +78,7 @@ class Session:
self.events: list[dict] = []
self._jpeg: Optional[bytes] = None
self._overlay: dict = {}
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)
@@ -92,6 +96,10 @@ class Session:
with self._lock:
return self._jpeg
def overlay(self) -> dict:
with self._lock:
return dict(self._overlay)
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.
@@ -119,6 +127,8 @@ class Session:
return {
"running": self._thread.is_alive(),
"source": self.source,
"whep_url": self.whep_url,
"preview": "webrtc" if self.whep_url else "mjpeg",
"model_path": self.model_path,
"error": self.error,
"frames": self.frames,
@@ -227,7 +237,11 @@ class Session:
self.fps = since / (now - tick)
tick, since = now, 0
self._render(frame, inside, outside, counter)
self._set_overlay(inside, outside, counter)
if not self.whep_url:
jpeg = live_render.render(self, frame, inside, outside, counter)
with self._lock:
self._jpeg = jpeg
except Exception as exc: # surfaced in status(), not swallowed
self.error = f"{type(exc).__name__}: {exc}"
finally:
@@ -263,59 +277,33 @@ class Session:
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)
def _set_overlay(self, detections, ignored, counter) -> None:
"""The same drawing `_render` burns into the JPEG, as geometry.
The browser draws it on a canvas over the WebRTC video, so the frames
themselves never pass through this process. Coordinates are in the
1280x720 space the pipeline works in; the canvas scales them.
"""
boxes = []
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()
state = "counted" if counter.counted_tracks.get(det.track_id) else "tracked"
boxes.append({"id": det.track_id, "b": [int(v) for v in det.bbox],
"c": round(float(det.confidence), 2), "s": state})
for det in ignored:
boxes.append({"id": det.track_id, "b": [int(v) for v in det.bbox],
"c": round(float(det.confidence), 2), "s": "ignored"})
payload = {
"frame": self.frames,
"line": {"y": self.line_y, "x_start": self.line_x_start,
"x_end": self.line_x_end, "margin": self.margin},
"boxes": boxes,
"loading": self.loading,
"unloading": self.unloading,
"net": self.loading - self.unloading,
"fps": round(self.fps, 1),
}
with self._lock:
self._overlay = payload
def _too_small(bbox, scale: float) -> bool:
"""`predict.py`'s perspective curve: a sack at the top of the frame is
@@ -339,74 +327,6 @@ def _too_small(bbox, scale: float) -> bool:
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
@@ -417,7 +337,8 @@ def start(source: str, model_path: str, line_y: int, line_x_start: int, line_x_e
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:
min_area_scale: float = 1.0, spatial_dedup: bool = False,
whep: str = "") -> dict:
global _session
with _guard:
if _session is not None and _session.status()["running"]:
@@ -429,7 +350,7 @@ def start(source: str, model_path: str, line_y: int, line_x_start: int, line_x_e
_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)
spatial_dedup, whep)
_session.start()
time.sleep(0.4) # let an immediate failure surface in the response
return _session.status()
@@ -459,9 +380,14 @@ def status() -> dict:
"fps": 0.0, "tracked": 0, "ignored": 0, "too_small": 0, "traced": 0,
"trace_path": "", "elapsed": 0.0, "events": [],
"error": "", "source": "", "model_path": "",
"whep_url": "", "preview": "mjpeg",
"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
def overlay() -> dict:
return _session.overlay() if _session is not None else {}