ariska/bytetrack-counter
1
0 Fork 0
forked from dsutanto/bytetrack-counter
Code Pull Requests Activity

Remove Clock Print

Browse Source
This commit is contained in:
proitlab committed 2026-06-25 18:59:38 +07:00
1 parent 93bc06b32b
commit a4fc55e15b
2 files changed
+727 -333
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

counter_live_rknn.py
+358 -167
View File
@@ -3,6 +3,7 @@ Edge production live counter — RTSP + YOLO RKNN + line crossing.
Runs on RK3588 hardware with RKNN model (320×320 input). Runs on RK3588 hardware with RKNN model (320×320 input).
Replaces the Jetson/TensorRT variant. Replaces the Jetson/TensorRT variant.
""" """
import numpy as np import numpy as np
import cv2 import cv2
import csv import csv
@@ -13,68 +14,73 @@ from datetime import datetime
from pathlib import Path from pathlib import Path
from dotenv import load_dotenv from dotenv import load_dotenv
load_dotenv() load_dotenv()
from rknnlite.api import RKNNLite from rknnlite.api import RKNNLite
from batch_store import BatchStore from batch_store import BatchStore
# --- config (override via env / .env) --- # --- config (override via env / .env) ---
OUTPUT_DIR = os.getenv('OUTPUT_DIR', '/opt/jetson-counter') OUTPUT_DIR = os.getenv("OUTPUT_DIR", "/opt/jetson-counter")
DB_PATH = os.getenv('DB_PATH', f'{OUTPUT_DIR}/jetson_counter.db') DB_PATH = os.getenv("DB_PATH", f"{OUTPUT_DIR}/jetson_counter.db")
STATE_FILE = os.getenv('STATE_FILE', f'{OUTPUT_DIR}/current_batch.json') STATE_FILE = os.getenv("STATE_FILE", f"{OUTPUT_DIR}/current_batch.json")
SOURCE = os.getenv('SOURCE', 'rtsp://user:pass@192.168.0.100:554/stream1') SOURCE = os.getenv("SOURCE", "rtsp://user:pass@192.168.0.100:554/stream1")
MODEL_PATH = os.getenv('MODEL_PATH', '/opt/jetson-counter/yolo11n.rknn') MODEL_PATH = os.getenv("MODEL_PATH", "/opt/jetson-counter/yolo11n.rknn")
CAMERA_NAME = os.getenv('CAMERA_NAME', 'CC1') CAMERA_NAME = os.getenv("CAMERA_NAME", "CC1")
OBJECT_LABEL = os.getenv('OBJECT_LABEL', 'ayam-potong') OBJECT_LABEL = os.getenv("OBJECT_LABEL", "ayam-potong")
CLASS_AYAM = os.getenv('CLASS_AYAM', 'ayam') CLASS_AYAM = os.getenv("CLASS_AYAM", "ayam")
CLASS_TALENAN = os.getenv('CLASS_TALENAN', 'talenan') CLASS_TALENAN = os.getenv("CLASS_TALENAN", "talenan")
LINE_X = int(os.getenv('LINE_X')) if os.getenv('LINE_X') else None LINE_X = int(os.getenv("LINE_X")) if os.getenv("LINE_X") else None
LINE_X_FRAC = float(os.getenv('LINE_X_FRAC', '0.5')) LINE_X_FRAC = float(os.getenv("LINE_X_FRAC", "0.5"))
CROSS_DIRECTION = os.getenv('CROSS_DIRECTION', 'rtl').lower() CROSS_DIRECTION = os.getenv("CROSS_DIRECTION", "rtl").lower()
IMGSZ = int(os.getenv('IMGSZ', '320')) IMGSZ = int(os.getenv("IMGSZ", "320"))
HALF = os.getenv('HALF', 'false').lower() == 'true' HALF = os.getenv("HALF", "false").lower() == "true"
CONF = float(os.getenv('CONF', '0.3')) CONF = float(os.getenv("CONF", "0.3"))
DEVICE = int(os.getenv('DEVICE', '0')) DEVICE = int(os.getenv("DEVICE", "0"))
# RKNN-specific — core mask for NPU # RKNN-specific — core mask for NPU
# 1 = core0, 2 = core1, 3 = core0+core1 (dual), 7 = all three # 1 = core0, 2 = core1, 3 = core0+core1 (dual), 7 = all three
CORE_MASK = int(os.getenv('CORE_MASK', '1')) CORE_MASK = int(os.getenv("CORE_MASK", "1"))
# YOLO decoder config # YOLO decoder config
NUM_CLASSES = int(os.getenv('NUM_CLASSES', '2')) NUM_CLASSES = int(os.getenv("NUM_CLASSES", "2"))
SCORE_SIGMOID = os.getenv('SCORE_SIGMOID', 'false').lower() == 'true' SCORE_SIGMOID = os.getenv("SCORE_SIGMOID", "false").lower() == "true"
DAILY_CUTOFF_TIME = os.getenv('DAILY_CUTOFF_TIME', '20:00') DAILY_CUTOFF_TIME = os.getenv("DAILY_CUTOFF_TIME", "20:00")
BATCH_TIMEOUT_SECONDS = float(os.getenv('BATCH_TIMEOUT_SECONDS', '300')) BATCH_TIMEOUT_SECONDS = float(os.getenv("BATCH_TIMEOUT_SECONDS", "300"))
IGNORE_BATCH_LABEL_TIMEOUT = float(os.getenv('IGNORE_BATCH_LABEL_TIMEOUT_SECONDS', '30')) IGNORE_BATCH_LABEL_TIMEOUT = float(
MIN_OBJECT_PER_BATCH = int(os.getenv('MIN_OBJECT_PER_BATCH', '60')) os.getenv("IGNORE_BATCH_LABEL_TIMEOUT_SECONDS", "30")
MIN_DURATION_PER_BATCH = int(os.getenv('MIN_DURATION_PER_BATCH', '60')) )
MIN_OBJECT_PER_BATCH = int(os.getenv("MIN_OBJECT_PER_BATCH", "60"))
MIN_DURATION_PER_BATCH = int(os.getenv("MIN_DURATION_PER_BATCH", "60"))
EXPORT_CSV = os.getenv('EXPORT_CSV', 'true').lower() == 'true' EXPORT_CSV = os.getenv("EXPORT_CSV", "true").lower() == "true"
CROSS_CSV = os.getenv('CROSS_CSV', f'{OUTPUT_DIR}/batch_crossings.csv') CROSS_CSV = os.getenv("CROSS_CSV", f"{OUTPUT_DIR}/batch_crossings.csv")
WARMUP_FRAMES = int(os.getenv('WARMUP_FRAMES', '30')) WARMUP_FRAMES = int(os.getenv("WARMUP_FRAMES", "30"))
RECONNECT_DELAY_SEC = int(os.getenv('RECONNECT_DELAY_SEC', '3')) RECONNECT_DELAY_SEC = int(os.getenv("RECONNECT_DELAY_SEC", "3"))
MAX_RECONNECT_ATTEMPTS = int(os.getenv('MAX_RECONNECT_ATTEMPTS', '0')) MAX_RECONNECT_ATTEMPTS = int(os.getenv("MAX_RECONNECT_ATTEMPTS", "0"))
FLUSH_EVERY_N_FRAMES = int(os.getenv('FLUSH_EVERY_N_FRAMES', '100')) FLUSH_EVERY_N_FRAMES = int(os.getenv("FLUSH_EVERY_N_FRAMES", "100"))
TRACKED_PRUNE_SEC = int(os.getenv('TRACKED_PRUNE_SEC', '300')) TRACKED_PRUNE_SEC = int(os.getenv("TRACKED_PRUNE_SEC", "300"))
RECORD_VIDEO = os.getenv('RECORD_VIDEO', 'false').lower() == 'true' RECORD_VIDEO = os.getenv("RECORD_VIDEO", "false").lower() == "true"
VIDEO_SEGMENT_SEC = int(os.getenv('VIDEO_SEGMENT_SEC', '3600')) VIDEO_SEGMENT_SEC = int(os.getenv("VIDEO_SEGMENT_SEC", "3600"))
OUTPUT_FPS = int(os.getenv('OUTPUT_FPS', '15')) OUTPUT_FPS = int(os.getenv("OUTPUT_FPS", "15"))
LIVE_STREAM_ENABLED = os.getenv('LIVE_STREAM_ENABLED', 'false').lower() == 'true' LIVE_STREAM_ENABLED = os.getenv("LIVE_STREAM_ENABLED", "false").lower() == "true"
LIVE_STREAM_FRAME_PATH = os.getenv('LIVE_STREAM_FRAME_PATH', '/dev/shm/jetson-counter/live_frame.jpg') LIVE_STREAM_FRAME_PATH = os.getenv(
LIVE_STREAM_QUALITY = int(os.getenv('LIVE_STREAM_QUALITY', '75')) "LIVE_STREAM_FRAME_PATH", "/dev/shm/jetson-counter/live_frame.jpg"
LIVE_STREAM_EVERY_N = int(os.getenv('LIVE_STREAM_EVERY_N', '2')) )
LIVE_STREAM_QUALITY = int(os.getenv("LIVE_STREAM_QUALITY", "75"))
LIVE_STREAM_EVERY_N = int(os.getenv("LIVE_STREAM_EVERY_N", "2"))
RTSP_FFMPEG_OPTIONS = os.getenv( RTSP_FFMPEG_OPTIONS = os.getenv(
'OPENCV_FFMPEG_CAPTURE_OPTIONS', "OPENCV_FFMPEG_CAPTURE_OPTIONS",
'rtsp_transport;tcp|fflags;nobuffer|flags;low_delay', "rtsp_transport;tcp|fflags;nobuffer|flags;low_delay",
) )
IS_LIVE = SOURCE.lower().startswith(('rtsp://', 'http://')) IS_LIVE = SOURCE.lower().startswith(("rtsp://", "http://"))
CROSS_FLASH_FRAMES = 12 CROSS_FLASH_FRAMES = 12
POPUP_LIFETIME = 20 POPUP_LIFETIME = 20
@@ -84,8 +90,14 @@ BATCH_PULSE_FRAMES = 20
SKELETON = [(0, 1), (4, 3), (1, 2), (3, 2), (2, 6), (2, 5), (2, 7), (7, 8)] SKELETON = [(0, 1), (4, 3), (1, 2), (3, 2), (2, 6), (2, 5), (2, 7), (7, 8)]
SK_COLORS = [ SK_COLORS = [
(0, 255, 255), (0, 255, 255), (255, 0, 255), (255, 0, 255), (0, 255, 255),
(0, 255, 0), (255, 255, 0), (0, 0, 255), (200, 200, 0), (0, 255, 255),
(255, 0, 255),
(255, 0, 255),
(0, 255, 0),
(255, 255, 0),
(0, 0, 255),
(200, 200, 0),
] ]
C_PANEL = (28, 24, 18) C_PANEL = (28, 24, 18)
@@ -105,7 +117,7 @@ shutdown_requested = False
def request_shutdown(signum, frame): def request_shutdown(signum, frame):
global shutdown_requested global shutdown_requested
shutdown_requested = True shutdown_requested = True
print('\nShutdown requested — finishing current frame...') print("\nShutdown requested — finishing current frame...")
signal.signal(signal.SIGINT, request_shutdown) signal.signal(signal.SIGINT, request_shutdown)
@@ -116,6 +128,7 @@ signal.signal(signal.SIGTERM, request_shutdown)
# YOLO output decoder (NMS only — boxes are pre-decoded by the model) # YOLO output decoder (NMS only — boxes are pre-decoded by the model)
# ============================================================================= # =============================================================================
def _nms(boxes, scores, iou_thr=0.45): def _nms(boxes, scores, iou_thr=0.45):
order = np.argsort(scores)[::-1] order = np.argsort(scores)[::-1]
keep = [] keep = []
@@ -132,7 +145,9 @@ def _nms(boxes, scores, iou_thr=0.45):
h = np.maximum(0.0, yy2 - yy1) h = np.maximum(0.0, yy2 - yy1)
inter = w * h inter = w * h
area_i = (boxes[idx, 2] - boxes[idx, 0]) * (boxes[idx, 3] - boxes[idx, 1]) area_i = (boxes[idx, 2] - boxes[idx, 0]) * (boxes[idx, 3] - boxes[idx, 1])
area_o = (boxes[order[1:], 2] - boxes[order[1:], 0]) * (boxes[order[1:], 3] - boxes[order[1:], 1]) area_o = (boxes[order[1:], 2] - boxes[order[1:], 0]) * (
boxes[order[1:], 3] - boxes[order[1:], 1]
)
iou = inter / (area_i + area_o - inter + 1e-16) iou = inter / (area_i + area_o - inter + 1e-16)
order = order[1:][iou < iou_thr] order = order[1:][iou < iou_thr]
return np.array(keep) return np.array(keep)
@@ -164,12 +179,13 @@ def _compute_iou(box1, boxes2):
# Simple IoU tracker (replaces bytetrack — same persist behaviour) # Simple IoU tracker (replaces bytetrack — same persist behaviour)
# ============================================================================= # =============================================================================
class SimpleTracker: class SimpleTracker:
def __init__(self, max_age=30, min_hits=1, iou_threshold=0.3): def __init__(self, max_age=30, min_hits=1, iou_threshold=0.3):
self.max_age = max_age self.max_age = max_age
self.min_hits = min_hits self.min_hits = min_hits
self.iou_threshold = iou_threshold self.iou_threshold = iou_threshold
self.tracks = {} # track_id -> {box, cx, age, hits, time_since_update} self.tracks = {} # track_id -> {box, cx, age, hits, time_since_update}
self.next_id = 1 self.next_id = 1
def update(self, detections): def update(self, detections):
@@ -177,33 +193,41 @@ class SimpleTracker:
now = time.monotonic() now = time.monotonic()
for tid in self.tracks: for tid in self.tracks:
self.tracks[tid]['time_since_update'] += 1 self.tracks[tid]["time_since_update"] += 1
matched_det = set() matched_det = set()
matched_track = set() matched_track = set()
assignments = [] # (track_id, det_idx) assignments = [] # (track_id, det_idx)
det_to_track = {} # det_idx → track_id det_to_track = {} # det_idx → track_id
if detections and self.tracks: if detections and self.tracks:
track_ids = list(self.tracks.keys()) track_ids = list(self.tracks.keys())
track_boxes = np.stack([self.tracks[t]['box'] for t in track_ids], axis=0) track_boxes = np.stack([self.tracks[t]["box"] for t in track_ids], axis=0)
for di, det in enumerate(detections): for di, det in enumerate(detections):
_, det_box = det _, det_box = det
ious = np.array([_compute_iou(det_box, track_boxes[t:t + 1]) for t in range(len(track_ids))]) ious = np.array(
[
_compute_iou(det_box, track_boxes[t : t + 1])
for t in range(len(track_ids))
]
)
best_j = int(np.argmax(ious)) best_j = int(np.argmax(ious))
if ious[best_j] >= self.iou_threshold and track_ids[best_j] not in matched_track: if (
ious[best_j] >= self.iou_threshold
and track_ids[best_j] not in matched_track
):
assignments.append((track_ids[best_j], di)) assignments.append((track_ids[best_j], di))
matched_track.add(track_ids[best_j]) matched_track.add(track_ids[best_j])
matched_det.add(di) matched_det.add(di)
for tid, di in assignments: for tid, di in assignments:
cx, box = detections[di] cx, box = detections[di]
self.tracks[tid]['cx'] = cx self.tracks[tid]["cx"] = cx
self.tracks[tid]['box'] = box self.tracks[tid]["box"] = box
self.tracks[tid]['hits'] += 1 self.tracks[tid]["hits"] += 1
self.tracks[tid]['time_since_update'] = 0 self.tracks[tid]["time_since_update"] = 0
self.tracks[tid]['last_update'] = now self.tracks[tid]["last_update"] = now
det_to_track[di] = tid det_to_track[di] = tid
for di, det in enumerate(detections): for di, det in enumerate(detections):
@@ -212,19 +236,27 @@ class SimpleTracker:
new_id = self.next_id new_id = self.next_id
self.next_id += 1 self.next_id += 1
self.tracks[new_id] = { self.tracks[new_id] = {
'cx': cx, 'box': box, 'hits': 1, "cx": cx,
'time_since_update': 0, 'last_update': now, "box": box,
"hits": 1,
"time_since_update": 0,
"last_update": now,
} }
det_to_track[di] = new_id det_to_track[di] = new_id
stale = [tid for tid, t in self.tracks.items() stale = [
if t['time_since_update'] > self.max_age] tid
for tid, t in self.tracks.items()
if t["time_since_update"] > self.max_age
]
for tid in stale: for tid in stale:
del self.tracks[tid] del self.tracks[tid]
track_map = {tid: self.tracks[tid]['cx'] track_map = {
for tid in self.tracks tid: self.tracks[tid]["cx"]
if self.tracks[tid]['hits'] >= self.min_hits} for tid in self.tracks
if self.tracks[tid]["hits"] >= self.min_hits
}
return track_map, det_to_track return track_map, det_to_track
@@ -232,9 +264,19 @@ class SimpleTracker:
# RKNN YOLO wrapper (detect output format: (1, 4+num_classes, N)) # RKNN YOLO wrapper (detect output format: (1, 4+num_classes, N))
# ============================================================================= # =============================================================================
class RKNNYOLO: class RKNNYOLO:
def __init__(self, model_path, core_mask=1, imgsz=320, conf=0.3, iou=0.45, def __init__(
num_classes=2, num_keypoints=0, score_sigmoid=False): self,
model_path,
core_mask=1,
imgsz=320,
conf=0.3,
iou=0.45,
num_classes=2,
num_keypoints=0,
score_sigmoid=False,
):
self.imgsz = imgsz self.imgsz = imgsz
self.conf = conf self.conf = conf
self.iou = iou self.iou = iou
@@ -245,19 +287,20 @@ class RKNNYOLO:
self.rknn = RKNNLite(verbose=False) self.rknn = RKNNLite(verbose=False)
ret = self.rknn.load_rknn(model_path) ret = self.rknn.load_rknn(model_path)
if ret != 0: if ret != 0:
raise RuntimeError(f'Failed to load RKNN model: {model_path}') raise RuntimeError(f"Failed to load RKNN model: {model_path}")
ret = self.rknn.init_runtime(core_mask=core_mask) ret = self.rknn.init_runtime(core_mask=core_mask)
if ret != 0: if ret != 0:
raise RuntimeError(f'Failed to init RKNN runtime (core_mask={core_mask})') raise RuntimeError(f"Failed to init RKNN runtime (core_mask={core_mask})")
try: try:
from rknnlite.api import RKNNLite as _RK from rknnlite.api import RKNNLite as _RK
sdk_ver = self.rknn.get_sdk_version() sdk_ver = self.rknn.get_sdk_version()
print(f'RKNN SDK version: {sdk_ver}') print(f"RKNN SDK version: {sdk_ver}")
except Exception: except Exception:
pass pass
print(f'RKNN model loaded: {model_path} imgsz={imgsz} core_mask={core_mask}') print(f"RKNN model loaded: {model_path} imgsz={imgsz} core_mask={core_mask}")
def _preprocess(self, frame): def _preprocess(self, frame):
"""Letterbox-resize to imgsz×imgsz, maintain aspect ratio, BGR→RGB, normalize.""" """Letterbox-resize to imgsz×imgsz, maintain aspect ratio, BGR→RGB, normalize."""
@@ -269,7 +312,7 @@ class RKNNYOLO:
letterbox = np.full((self.imgsz, self.imgsz, 3), 114, dtype=np.uint8) letterbox = np.full((self.imgsz, self.imgsz, 3), 114, dtype=np.uint8)
dy = (self.imgsz - nh) // 2 dy = (self.imgsz - nh) // 2
dx = (self.imgsz - nw) // 2 dx = (self.imgsz - nw) // 2
letterbox[dy:dy + nh, dx:dx + nw] = resized letterbox[dy : dy + nh, dx : dx + nw] = resized
rgb = cv2.cvtColor(letterbox, cv2.COLOR_BGR2RGB) rgb = cv2.cvtColor(letterbox, cv2.COLOR_BGR2RGB)
gains = np.array([scale, scale, dy, dx], dtype=np.float32) gains = np.array([scale, scale, dy, dx], dtype=np.float32)
@@ -295,7 +338,7 @@ class RKNNYOLO:
if out.shape[0] == self.num_classes + 4: if out.shape[0] == self.num_classes + 4:
out = out.T # (C, N) → (N, C) out = out.T # (C, N) → (N, C)
boxes_cxcywh = out[:, :4].copy() # cx, cy, w, h at model resolution boxes_cxcywh = out[:, :4].copy() # cx, cy, w, h at model resolution
cls_raw = out[:, 4:].copy() cls_raw = out[:, 4:].copy()
if self.score_sigmoid: if self.score_sigmoid:
@@ -303,12 +346,15 @@ class RKNNYOLO:
else: else:
cls_scores = cls_raw cls_scores = cls_raw
boxes_xyxy = np.stack([ boxes_xyxy = np.stack(
boxes_cxcywh[:, 0] - boxes_cxcywh[:, 2] / 2, [
boxes_cxcywh[:, 1] - boxes_cxcywh[:, 3] / 2, boxes_cxcywh[:, 0] - boxes_cxcywh[:, 2] / 2,
boxes_cxcywh[:, 0] + boxes_cxcywh[:, 2] / 2, boxes_cxcywh[:, 1] - boxes_cxcywh[:, 3] / 2,
boxes_cxcywh[:, 1] + boxes_cxcywh[:, 3] / 2, boxes_cxcywh[:, 0] + boxes_cxcywh[:, 2] / 2,
], axis=1) boxes_cxcywh[:, 1] + boxes_cxcywh[:, 3] / 2,
],
axis=1,
)
max_scores = cls_scores.max(axis=1) max_scores = cls_scores.max(axis=1)
class_ids = cls_scores.argmax(axis=1) class_ids = cls_scores.argmax(axis=1)
@@ -339,12 +385,14 @@ class RKNNYOLO:
keep = _nms(bboxes[idx], scores[idx], iou_thr=self.iou) keep = _nms(bboxes[idx], scores[idx], iou_thr=self.iou)
for k in keep: for k in keep:
j = idx[k] j = idx[k]
detections.append({ detections.append(
'bbox': bboxes[j].tolist(), {
'score': float(scores[j]), "bbox": bboxes[j].tolist(),
'cls': int(clses[j]), "score": float(scores[j]),
'keypoints': None, "cls": int(clses[j]),
}) "keypoints": None,
}
)
return detections return detections
def release(self): def release(self):
@@ -365,41 +413,41 @@ def resolve_line_x(frame_width):
def crossed_line(prev_cx, cx, line_x, direction=CROSS_DIRECTION): def crossed_line(prev_cx, cx, line_x, direction=CROSS_DIRECTION):
if direction == 'ltr': if direction == "ltr":
return prev_cx < line_x <= cx return prev_cx < line_x <= cx
if direction == 'both': if direction == "both":
return (prev_cx > line_x >= cx) or (prev_cx < line_x <= cx) return (prev_cx > line_x >= cx) or (prev_cx < line_x <= cx)
return prev_cx > line_x >= cx return prev_cx > line_x >= cx
def now_str(): def now_str():
return datetime.now().strftime('%Y-%m-%d %H:%M:%S') return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
def open_capture(source): def open_capture(source):
if source.lower().startswith(('rtsp://', 'http://')): if source.lower().startswith(("rtsp://", "http://")):
os.environ['OPENCV_FFMPEG_CAPTURE_OPTIONS'] = RTSP_FFMPEG_OPTIONS os.environ["OPENCV_FFMPEG_CAPTURE_OPTIONS"] = RTSP_FFMPEG_OPTIONS
cap = cv2.VideoCapture(source, cv2.CAP_FFMPEG) cap = cv2.VideoCapture(source, cv2.CAP_FFMPEG)
cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
return cap return cap
def warmup_stream(cap, n=WARMUP_FRAMES): def warmup_stream(cap, n=WARMUP_FRAMES):
print('Warming up stream...') print("Warming up stream...")
for _ in range(n): for _ in range(n):
cap.read() cap.read()
print('Stream ready!') print("Stream ready!")
def open_video_writer(path, w, h, fps): def open_video_writer(path, w, h, fps):
return cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*'avc1'), fps, (w, h)) return cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*"avc1"), fps, (w, h))
class CsvLogger: class CsvLogger:
def __init__(self, path, header): def __init__(self, path, header):
Path(path).parent.mkdir(parents=True, exist_ok=True) Path(path).parent.mkdir(parents=True, exist_ok=True)
new_file = not Path(path).exists() or Path(path).stat().st_size == 0 new_file = not Path(path).exists() or Path(path).stat().st_size == 0
self.file = open(path, 'a', newline='', buffering=1) self.file = open(path, "a", newline="", buffering=1)
self.writer = csv.writer(self.file) self.writer = csv.writer(self.file)
if new_file: if new_file:
self.writer.writerow(header) self.writer.writerow(header)
@@ -424,8 +472,8 @@ class VideoSegmentWriter:
self._open_next() self._open_next()
def _segment_path(self): def _segment_path(self):
ts = datetime.now().strftime('%Y%m%d_%H%M%S') ts = datetime.now().strftime("%Y%m%d_%H%M%S")
return str(self.output_dir / f'live_{ts}.mp4') return str(self.output_dir / f"live_{ts}.mp4")
def _open_next(self): def _open_next(self):
if self.writer is not None: if self.writer is not None:
@@ -433,7 +481,7 @@ class VideoSegmentWriter:
path = self._segment_path() path = self._segment_path()
self.writer = open_video_writer(path, self.w, self.h, self.fps) self.writer = open_video_writer(path, self.w, self.h, self.fps)
self.segment_start = time.monotonic() self.segment_start = time.monotonic()
print(f'Recording segment: {path}') print(f"Recording segment: {path}")
def write(self, frame): def write(self, frame):
if time.monotonic() - self.segment_start >= self.segment_sec: if time.monotonic() - self.segment_start >= self.segment_sec:
@@ -446,7 +494,9 @@ class VideoSegmentWriter:
def prune_stale_tracks(tracked, now_mono): def prune_stale_tracks(tracked, now_mono):
stale = [tid for tid, (_, ts) in tracked.items() if now_mono - ts > TRACKED_PRUNE_SEC] stale = [
tid for tid, (_, ts) in tracked.items() if now_mono - ts > TRACKED_PRUNE_SEC
]
for tid in stale: for tid in stale:
del tracked[tid] del tracked[tid]
@@ -484,7 +534,16 @@ def draw_elegant_counting_line(img, line_x, h, pulse_remaining=0):
y_end = min(y + dash_len, h) y_end = min(y + dash_len, h)
cv2.line(img, (line_x, y), (line_x, y_end), C_LINE_CORE, 2, cv2.LINE_AA) cv2.line(img, (line_x, y), (line_x, y_end), C_LINE_CORE, 2, cv2.LINE_AA)
y += dash_len + gap y += dash_len + gap
cv2.putText(img, 'COUNT LINE', (line_x - 46, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.42, C_LINE_CORE, 1, cv2.LINE_AA) cv2.putText(
img,
"COUNT LINE",
(line_x - 46, 24),
cv2.FONT_HERSHEY_SIMPLEX,
0.42,
C_LINE_CORE,
1,
cv2.LINE_AA,
)
def draw_hero_count(img, line_x, h, count, pulse_remaining=0): def draw_hero_count(img, line_x, h, count, pulse_remaining=0):
@@ -495,34 +554,120 @@ def draw_hero_count(img, line_x, h, count, pulse_remaining=0):
(tw, th), _ = cv2.getTextSize(text, font, font_scale, thickness) (tw, th), _ = cv2.getTextSize(text, font, font_scale, thickness)
pad = 14 pad = 14
tx, ty = line_x - tw // 2, h // 2 + th // 2 tx, ty = line_x - tw // 2, h // 2 + th // 2
overlay_rect(img, tx - pad, ty - th - pad, tx + tw + pad, ty + pad // 2, C_PANEL, alpha=0.78) overlay_rect(
cv2.rectangle(img, (tx - pad, ty - th - pad), (tx + tw + pad, ty + pad // 2), C_LINE_CORE, 2) img, tx - pad, ty - th - pad, tx + tw + pad, ty + pad // 2, C_PANEL, alpha=0.78
)
cv2.rectangle(
img, (tx - pad, ty - th - pad), (tx + tw + pad, ty + pad // 2), C_LINE_CORE, 2
)
cv2.putText(img, text, (tx, ty), font, font_scale, C_GREEN, thickness, cv2.LINE_AA) cv2.putText(img, text, (tx, ty), font, font_scale, C_GREEN, thickness, cv2.LINE_AA)
def draw_hud(img, w, batch_num, batch_count, total_ayam, elapsed_sec, rate, camera_id, clock): def draw_hud(
img, w, batch_num, batch_count, total_ayam, elapsed_sec, rate, camera_id, clock
):
bar_h = 52 bar_h = 52
overlay_rect(img, 0, 0, w, bar_h, C_PANEL, alpha=0.72) overlay_rect(img, 0, 0, w, bar_h, C_PANEL, alpha=0.72)
cv2.line(img, (0, bar_h), (w, bar_h), C_BORDER, 1) cv2.line(img, (0, bar_h), (w, bar_h), C_BORDER, 1)
cv2.putText(img, 'BATCH', (16, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) cv2.putText(
batch_label = str(batch_num) if batch_num else '—' img, "BATCH", (16, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA
cv2.putText(img, batch_label, (16, 44), cv2.FONT_HERSHEY_SIMPLEX, 0.9, C_ACCENT, 2, cv2.LINE_AA) )
cv2.putText(img, 'COUNT', (100, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) batch_label = str(batch_num) if batch_num else "—"
cv2.putText(img, str(batch_count), (100, 44), cv2.FONT_HERSHEY_SIMPLEX, 0.9, C_GREEN, 2, cv2.LINE_AA) cv2.putText(
cv2.putText(img, 'TOTAL', (190, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) img,
cv2.putText(img, str(total_ayam), (190, 44), cv2.FONT_HERSHEY_SIMPLEX, 0.7, C_TEXT, 1, cv2.LINE_AA) batch_label,
cv2.putText(img, 'UPTIME', (280, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) (16, 44),
cv2.putText(img, f'{elapsed_sec / 3600:.1f}h', (280, 44), cv2.FONT_HERSHEY_SIMPLEX, 0.7, C_TEXT, 1, cv2.LINE_AA) cv2.FONT_HERSHEY_SIMPLEX,
cv2.putText(img, 'RATE', (380, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) 0.9,
cv2.putText(img, f'{rate:.1f}/min', (380, 44), cv2.FONT_HERSHEY_SIMPLEX, 0.7, C_ACCENT, 1, cv2.LINE_AA) C_ACCENT,
cv2.putText(img, clock, (w - 180, 36), cv2.FONT_HERSHEY_SIMPLEX, 0.55, C_TEXT, 1, cv2.LINE_AA) 2,
cv2.putText(img, f'CAM {camera_id}', (w - 180, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) cv2.LINE_AA,
)
cv2.putText(
img, "COUNT", (100, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA
)
cv2.putText(
img,
str(batch_count),
(100, 44),
cv2.FONT_HERSHEY_SIMPLEX,
0.9,
C_GREEN,
2,
cv2.LINE_AA,
)
cv2.putText(
img, "TOTAL", (190, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA
)
cv2.putText(
img,
str(total_ayam),
(190, 44),
cv2.FONT_HERSHEY_SIMPLEX,
0.7,
C_TEXT,
1,
cv2.LINE_AA,
)
cv2.putText(
img,
"UPTIME",
(280, 20),
cv2.FONT_HERSHEY_SIMPLEX,
0.45,
C_MUTED,
1,
cv2.LINE_AA,
)
cv2.putText(
img,
f"{elapsed_sec / 3600:.1f}h",
(280, 44),
cv2.FONT_HERSHEY_SIMPLEX,
0.7,
C_TEXT,
1,
cv2.LINE_AA,
)
cv2.putText(
img, "RATE", (380, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA
)
cv2.putText(
img,
f"{rate:.1f}/min",
(380, 44),
cv2.FONT_HERSHEY_SIMPLEX,
0.7,
C_ACCENT,
1,
cv2.LINE_AA,
)
# cv2.putText(img, clock, (w - 180, 36), cv2.FONT_HERSHEY_SIMPLEX, 0.55, C_TEXT, 1, cv2.LINE_AA)
cv2.putText(
img,
f"CAM {camera_id}",
(w - 180, 20),
cv2.FONT_HERSHEY_SIMPLEX,
0.45,
C_MUTED,
1,
cv2.LINE_AA,
)
def draw_footer(img, w, h, frame_idx, live_tag): def draw_footer(img, w, h, frame_idx, live_tag):
bar_h = 28 bar_h = 28
overlay_rect(img, 0, h - bar_h, w, h, C_PANEL, alpha=0.55) overlay_rect(img, 0, h - bar_h, w, h, C_PANEL, alpha=0.55)
cv2.putText(img, f'{live_tag} | Frame {frame_idx}', (12, h - 9), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) cv2.putText(
img,
f"{live_tag} | Frame {frame_idx}",
(12, h - 9),
cv2.FONT_HERSHEY_SIMPLEX,
0.45,
C_MUTED,
1,
cv2.LINE_AA,
)
def draw_skeleton_bold(img, kpts): def draw_skeleton_bold(img, kpts):
@@ -542,28 +687,39 @@ def draw_skeleton_bold(img, kpts):
def draw_popups(img, popups, frame_idx): def draw_popups(img, popups, frame_idx):
alive = [] alive = []
for pop in popups: for pop in popups:
age = frame_idx - pop['born'] age = frame_idx - pop["born"]
if age > POPUP_LIFETIME: if age > POPUP_LIFETIME:
continue continue
alive.append(pop) alive.append(pop)
fade = 1.0 - age / POPUP_LIFETIME fade = 1.0 - age / POPUP_LIFETIME
y = pop['y'] - int(age * 1.8) y = pop["y"] - int(age * 1.8)
color = (int(C_GREEN[0] * fade), int(C_GREEN[1] * fade), int(C_GREEN[2] * fade)) color = (int(C_GREEN[0] * fade), int(C_GREEN[1] * fade), int(C_GREEN[2] * fade))
cv2.putText(img, pop['text'], (pop['x'], y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2, cv2.LINE_AA) cv2.putText(
img,
pop["text"],
(pop["x"], y),
cv2.FONT_HERSHEY_SIMPLEX,
0.7,
color,
2,
cv2.LINE_AA,
)
return alive return alive
def draw_batch_banner(img, w, batch_num, pulse_remaining): def draw_batch_banner(img, w, batch_num, pulse_remaining):
if pulse_remaining <= 0: if pulse_remaining <= 0:
return return
text = f'NEW BATCH {batch_num}' text = f"NEW BATCH {batch_num}"
font = cv2.FONT_HERSHEY_SIMPLEX font = cv2.FONT_HERSHEY_SIMPLEX
(tw, th), _ = cv2.getTextSize(text, font, 0.8, 2) (tw, th), _ = cv2.getTextSize(text, font, 0.8, 2)
x1, y1 = w // 2 - tw // 2 - 16, 62 x1, y1 = w // 2 - tw // 2 - 16, 62
x2, y2 = w // 2 + tw // 2 + 16, 62 + th + 20 x2, y2 = w // 2 + tw // 2 + 16, 62 + th + 20
overlay_rect(img, x1, y1, x2, y2, C_PANEL, alpha=0.7) overlay_rect(img, x1, y1, x2, y2, C_PANEL, alpha=0.7)
cv2.rectangle(img, (x1, y1), (x2, y2), C_ACCENT, 2) cv2.rectangle(img, (x1, y1), (x2, y2), C_ACCENT, 2)
cv2.putText(img, text, (w // 2 - tw // 2, 62 + th + 4), font, 0.8, C_ACCENT, 2, cv2.LINE_AA) cv2.putText(
img, text, (w // 2 - tw // 2, 62 + th + 4), font, 0.8, C_ACCENT, 2, cv2.LINE_AA
)
def connect_stream(source, warmup=WARMUP_FRAMES): def connect_stream(source, warmup=WARMUP_FRAMES):
@@ -573,11 +729,13 @@ def connect_stream(source, warmup=WARMUP_FRAMES):
if not cap.isOpened(): if not cap.isOpened():
attempts += 1 attempts += 1
if MAX_RECONNECT_ATTEMPTS and attempts >= MAX_RECONNECT_ATTEMPTS: if MAX_RECONNECT_ATTEMPTS and attempts >= MAX_RECONNECT_ATTEMPTS:
raise RuntimeError(f'Cannot open source after {attempts} attempts: {source}') raise RuntimeError(
print(f'Cannot open source, retry in {RECONNECT_DELAY_SEC}s...') f"Cannot open source after {attempts} attempts: {source}"
)
print(f"Cannot open source, retry in {RECONNECT_DELAY_SEC}s...")
time.sleep(RECONNECT_DELAY_SEC) time.sleep(RECONNECT_DELAY_SEC)
continue continue
if warmup > 0 and source.lower().startswith(('rtsp://', 'http://')): if warmup > 0 and source.lower().startswith(("rtsp://", "http://")):
warmup_stream(cap, warmup) warmup_stream(cap, warmup)
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
@@ -592,6 +750,7 @@ def connect_stream(source, warmup=WARMUP_FRAMES):
# Main loop # Main loop
# ============================================================================= # =============================================================================
def run(): def run():
global shutdown_requested global shutdown_requested
@@ -605,13 +764,15 @@ def run():
ignore_batch_label_timeout=IGNORE_BATCH_LABEL_TIMEOUT, ignore_batch_label_timeout=IGNORE_BATCH_LABEL_TIMEOUT,
min_object_per_batch=MIN_OBJECT_PER_BATCH, min_object_per_batch=MIN_OBJECT_PER_BATCH,
min_duration_per_batch=MIN_DURATION_PER_BATCH, min_duration_per_batch=MIN_DURATION_PER_BATCH,
logger=lambda msg: print(f'[{now_str()}] {msg}'), logger=lambda msg: print(f"[{now_str()}] {msg}"),
) )
store.start_cutoff_watcher() store.start_cutoff_watcher()
cross_logger = None cross_logger = None
if EXPORT_CSV: if EXPORT_CSV:
cross_logger = CsvLogger(CROSS_CSV, ['batch', 'frame', 'timestamp', 'chicken_id']) cross_logger = CsvLogger(
CROSS_CSV, ["batch", "frame", "timestamp", "chicken_id"]
)
# Load RKNN model # Load RKNN model
model = RKNNYOLO( model = RKNNYOLO(
@@ -627,8 +788,8 @@ def run():
# class index 0 → CLASS_AYAM, index 1 → CLASS_TALENAN (or env-specified) # class index 0 → CLASS_AYAM, index 1 → CLASS_TALENAN (or env-specified)
# Use class names in env order: first CLASS_AYAM → id 0, then CLASS_TALENAN → id 1 # Use class names in env order: first CLASS_AYAM → id 0, then CLASS_TALENAN → id 1
CLASS_IDS = { CLASS_IDS = {
os.getenv('CLASS_AYAM', 'ayam'): 0, os.getenv("CLASS_AYAM", "ayam"): 0,
os.getenv('CLASS_TALENAN', 'talenan'): 1, os.getenv("CLASS_TALENAN", "talenan"): 1,
} }
ayam_cls = CLASS_IDS[CLASS_AYAM] ayam_cls = CLASS_IDS[CLASS_AYAM]
talenan_cls = CLASS_IDS[CLASS_TALENAN] talenan_cls = CLASS_IDS[CLASS_TALENAN]
@@ -655,10 +816,12 @@ def run():
return return
line_x = resolve_line_x(w) line_x = resolve_line_x(w)
print(f'RKNN counter | {w}x{h} @ {fps}fps | line x={line_x} | cross={CROSS_DIRECTION}') print(
print(f'Model: {MODEL_PATH} | imgsz={IMGSZ} | core_mask={CORE_MASK}') f"RKNN counter | {w}x{h} @ {fps}fps | line x={line_x} | cross={CROSS_DIRECTION}"
print(f'DB: {DB_PATH}') )
print(f'State: {STATE_FILE}') print(f"Model: {MODEL_PATH} | imgsz={IMGSZ} | core_mask={CORE_MASK}")
print(f"DB: {DB_PATH}")
print(f"State: {STATE_FILE}")
if RECORD_VIDEO: if RECORD_VIDEO:
video_writer = VideoSegmentWriter(OUTPUT_DIR, w, h, fps, VIDEO_SEGMENT_SEC) video_writer = VideoSegmentWriter(OUTPUT_DIR, w, h, fps, VIDEO_SEGMENT_SEC)
@@ -671,7 +834,9 @@ def run():
if not IS_LIVE: if not IS_LIVE:
break break
reconnect_count += 1 reconnect_count += 1
print(f'Stream dropped (attempt {reconnect_count}), reconnecting in {RECONNECT_DELAY_SEC}s...') print(
f"Stream dropped (attempt {reconnect_count}), reconnecting in {RECONNECT_DELAY_SEC}s..."
)
cap.release() cap.release()
time.sleep(RECONNECT_DELAY_SEC) time.sleep(RECONNECT_DELAY_SEC)
cap, w, h, fps = connect_stream(SOURCE) cap, w, h, fps = connect_stream(SOURCE)
@@ -691,12 +856,12 @@ def run():
if detections: if detections:
ayam_dets = [] # list of (cx, xywh_box) ayam_dets = [] # list of (cx, xywh_box)
talenan_dets = [] talenan_dets = []
ayam_kpts_map = {} # det_idx → keypoints ayam_kpts_map = {} # det_idx → keypoints
talenan_kpts_map = {} talenan_kpts_map = {}
for di, det in enumerate(detections): for di, det in enumerate(detections):
bbox = det['bbox'] bbox = det["bbox"]
cls_id = det['cls'] cls_id = det["cls"]
cx = (bbox[0] + bbox[2]) / 2.0 cx = (bbox[0] + bbox[2]) / 2.0
x1, y1, x2, y2 = bbox x1, y1, x2, y2 = bbox
wb, hb = x2 - x1, y2 - y1 wb, hb = x2 - x1, y2 - y1
@@ -704,12 +869,12 @@ def run():
if cls_id == talenan_cls: if cls_id == talenan_cls:
talenan_dets.append((cx, box_cxcywh)) talenan_dets.append((cx, box_cxcywh))
if det['keypoints'] is not None: if det["keypoints"] is not None:
talenan_kpts_map[len(talenan_dets) - 1] = det['keypoints'] talenan_kpts_map[len(talenan_dets) - 1] = det["keypoints"]
elif cls_id == ayam_cls: elif cls_id == ayam_cls:
ayam_dets.append((cx, box_cxcywh)) ayam_dets.append((cx, box_cxcywh))
if det['keypoints'] is not None: if det["keypoints"] is not None:
ayam_kpts_map[len(ayam_dets) - 1] = det['keypoints'] ayam_kpts_map[len(ayam_dets) - 1] = det["keypoints"]
# Track ayam — returns (track_id → cx, detection_idx → track_id) # Track ayam — returns (track_id → cx, detection_idx → track_id)
ayam_cx_map, ayam_det_to_track = ayam_tracker.update(ayam_dets) ayam_cx_map, ayam_det_to_track = ayam_tracker.update(ayam_dets)
@@ -725,17 +890,22 @@ def run():
if tid in talenan_tracker.tracks: if tid in talenan_tracker.tracks:
if tid in talenan_tracked: if tid in talenan_tracked:
prev_cx = talenan_tracked[tid][0] prev_cx = talenan_tracked[tid][0]
if crossed_line(prev_cx, cx, line_x) and tid not in talenan_line_crossed: if (
crossed_line(prev_cx, cx, line_x)
and tid not in talenan_line_crossed
):
talenan_line_crossed.add(tid) talenan_line_crossed.add(tid)
if store.record_talenan_crossing(tid): if store.record_talenan_crossing(tid):
batch_closed_frame = True batch_closed_frame = True
talenan_cross_flash[tid] = CROSS_FLASH_FRAMES talenan_cross_flash[tid] = CROSS_FLASH_FRAMES
popups.append({ popups.append(
'x': int(cx) - 20, {
'y': int(box[1]), "x": int(cx) - 20,
'born': frame_idx, "y": int(box[1]),
'text': 'BATCH CLOSED', "born": frame_idx,
}) "text": "BATCH CLOSED",
}
)
talenan_tracked[tid] = (cx, mono) talenan_tracked[tid] = (cx, mono)
# Process ayam crossings # Process ayam crossings
@@ -746,24 +916,33 @@ def run():
if tid in ayam_tracker.tracks: if tid in ayam_tracker.tracks:
if tid in ayam_tracked: if tid in ayam_tracked:
prev_cx = ayam_tracked[tid][0] prev_cx = ayam_tracked[tid][0]
if crossed_line(prev_cx, cx, line_x) and tid not in ayam_line_crossed: if (
crossed_line(prev_cx, cx, line_x)
and tid not in ayam_line_crossed
):
ayam_line_crossed.add(tid) ayam_line_crossed.add(tid)
_, started_new = store.record_ayam_crossing(tid) _, started_new = store.record_ayam_crossing(tid)
if cross_logger: if cross_logger:
cross_logger.write_row([ cross_logger.write_row(
store.current_batch_number, frame_idx, [
datetime.now().isoformat(), tid, store.current_batch_number,
]) frame_idx,
datetime.now().isoformat(),
tid,
]
)
ayam_crossed_frame = True ayam_crossed_frame = True
if started_new: if started_new:
batch_started_frame = True batch_started_frame = True
ayam_cross_flash[tid] = CROSS_FLASH_FRAMES ayam_cross_flash[tid] = CROSS_FLASH_FRAMES
popups.append({ popups.append(
'x': int(cx) - 12, {
'y': int(box[1]), "x": int(cx) - 12,
'born': frame_idx, "y": int(box[1]),
'text': '+1', "born": frame_idx,
}) "text": "+1",
}
)
ayam_tracked[tid] = (cx, mono) ayam_tracked[tid] = (cx, mono)
# Draw talenan # Draw talenan
@@ -778,7 +957,7 @@ def run():
flash = talenan_cross_flash.get(tid, 0) flash = talenan_cross_flash.get(tid, 0)
color = C_GREEN if flash > 0 else C_TALENAN_BOX color = C_GREEN if flash > 0 else C_TALENAN_BOX
cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3 if flash > 0 else 2) cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3 if flash > 0 else 2)
draw_pill(frame, f'TALENAN {tid}', x1, y1 - 4, color) draw_pill(frame, f"TALENAN {tid}", x1, y1 - 4, color)
# Draw ayam # Draw ayam
for di, (cx, box) in enumerate(ayam_dets): for di, (cx, box) in enumerate(ayam_dets):
@@ -792,7 +971,7 @@ def run():
flash = ayam_cross_flash.get(tid, 0) flash = ayam_cross_flash.get(tid, 0)
color = C_GREEN if flash > 0 else C_AYAM_BOX color = C_GREEN if flash > 0 else C_AYAM_BOX
cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3 if flash > 0 else 2) cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3 if flash > 0 else 2)
draw_pill(frame, f'ID {tid}', x1, y1 - 4, color) draw_pill(frame, f"ID {tid}", x1, y1 - 4, color)
kpts = ayam_kpts_map.get(di) kpts = ayam_kpts_map.get(di)
if kpts is not None: if kpts is not None:
draw_skeleton_bold(frame, kpts) draw_skeleton_bold(frame, kpts)
@@ -812,9 +991,19 @@ def run():
draw_elegant_counting_line(frame, line_x, h, line_pulse) draw_elegant_counting_line(frame, line_x, h, line_pulse)
draw_hero_count(frame, line_x, h, batch_count, count_pulse) draw_hero_count(frame, line_x, h, batch_count, count_pulse)
draw_hud(frame, w, batch_num, batch_count, display_total, elapsed, rate, CAMERA_NAME, now_str()) draw_hud(
frame,
w,
batch_num,
batch_count,
display_total,
elapsed,
rate,
CAMERA_NAME,
now_str(),
)
draw_batch_banner(frame, w, batch_num, batch_pulse) draw_batch_banner(frame, w, batch_num, batch_pulse)
draw_footer(frame, w, h, frame_idx, 'LIVE-RKNN' if IS_LIVE else 'FILE-RKNN') draw_footer(frame, w, h, frame_idx, "LIVE-RKNN" if IS_LIVE else "FILE-RKNN")
popups = draw_popups(frame, popups, frame_idx) popups = draw_popups(frame, popups, frame_idx)
for flash_store in (ayam_cross_flash, talenan_cross_flash): for flash_store in (ayam_cross_flash, talenan_cross_flash):
@@ -832,8 +1021,10 @@ def run():
if LIVE_STREAM_ENABLED and frame_idx % LIVE_STREAM_EVERY_N == 0: if LIVE_STREAM_ENABLED and frame_idx % LIVE_STREAM_EVERY_N == 0:
try: try:
Path(LIVE_STREAM_FRAME_PATH).parent.mkdir(parents=True, exist_ok=True) Path(LIVE_STREAM_FRAME_PATH).parent.mkdir(parents=True, exist_ok=True)
_, jpeg = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, LIVE_STREAM_QUALITY]) _, jpeg = cv2.imencode(
with open(LIVE_STREAM_FRAME_PATH, 'wb') as f: ".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, LIVE_STREAM_QUALITY]
)
with open(LIVE_STREAM_FRAME_PATH, "wb") as f:
f.write(jpeg.tobytes()) f.write(jpeg.tobytes())
except Exception: except Exception:
pass pass
@@ -841,8 +1032,8 @@ def run():
frame_idx += 1 frame_idx += 1
if frame_idx % FLUSH_EVERY_N_FRAMES == 0: if frame_idx % FLUSH_EVERY_N_FRAMES == 0:
print( print(
f'[{now_str()}] Frame {frame_idx} | Batch {batch_num}: {batch_count} ' f"[{now_str()}] Frame {frame_idx} | Batch {batch_num}: {batch_count} "
f'| Total: {display_total} | Uptime {elapsed / 3600:.2f}h' f"| Total: {display_total} | Uptime {elapsed / 3600:.2f}h"
) )
prune_stale_tracks(ayam_tracked, mono) prune_stale_tracks(ayam_tracked, mono)
prune_stale_tracks(talenan_tracked, mono) prune_stale_tracks(talenan_tracked, mono)
@@ -855,8 +1046,8 @@ def run():
model.release() model.release()
store.shutdown() store.shutdown()
print('\n=== Batch Summary (SQLite) ===') print("\n=== Batch Summary (SQLite) ===")
print(f'Database: {DB_PATH}') print(f"Database: {DB_PATH}")
# Tracked state dicts: track_id → (cx, monotonic_time) # Tracked state dicts: track_id → (cx, monotonic_time)
@@ -864,5 +1055,5 @@ ayam_tracked = {}
talenan_tracked = {} talenan_tracked = {}
if __name__ == '__main__': if __name__ == "__main__":
run() run()
counter_live_rknn_bytetrack.py
+369 -166
View File
@@ -3,6 +3,7 @@ Edge production live counter — RTSP + YOLO RKNN + ByteTrack + line crossing.
Runs on RK3588 hardware with RKNN model (320×320 input). Runs on RK3588 hardware with RKNN model (320×320 input).
Uses ByteTrack (Kalman filter + two-stage IoU association) for tracking. Uses ByteTrack (Kalman filter + two-stage IoU association) for tracking.
""" """
import numpy as np import numpy as np
import cv2 import cv2
import csv import csv
@@ -13,74 +14,79 @@ from datetime import datetime
from pathlib import Path from pathlib import Path
from dotenv import load_dotenv from dotenv import load_dotenv
load_dotenv() load_dotenv()
from rknnlite.api import RKNNLite from rknnlite.api import RKNNLite
from batch_store import BatchStore from batch_store import BatchStore
# --- config (override via env / .env) --- # --- config (override via env / .env) ---
OUTPUT_DIR = os.getenv('OUTPUT_DIR', '/opt/jetson-counter') OUTPUT_DIR = os.getenv("OUTPUT_DIR", "/opt/jetson-counter")
DB_PATH = os.getenv('DB_PATH', f'{OUTPUT_DIR}/jetson_counter.db') DB_PATH = os.getenv("DB_PATH", f"{OUTPUT_DIR}/jetson_counter.db")
STATE_FILE = os.getenv('STATE_FILE', f'{OUTPUT_DIR}/current_batch.json') STATE_FILE = os.getenv("STATE_FILE", f"{OUTPUT_DIR}/current_batch.json")
SOURCE = os.getenv('SOURCE', 'rtsp://user:pass@192.168.0.100:554/stream1') SOURCE = os.getenv("SOURCE", "rtsp://user:pass@192.168.0.100:554/stream1")
MODEL_PATH = os.getenv('MODEL_PATH', '/opt/jetson-counter/yolo11n.rknn') MODEL_PATH = os.getenv("MODEL_PATH", "/opt/jetson-counter/yolo11n.rknn")
CAMERA_NAME = os.getenv('CAMERA_NAME', 'CC1') CAMERA_NAME = os.getenv("CAMERA_NAME", "CC1")
OBJECT_LABEL = os.getenv('OBJECT_LABEL', 'ayam-potong') OBJECT_LABEL = os.getenv("OBJECT_LABEL", "ayam-potong")
CLASS_AYAM = os.getenv('CLASS_AYAM', 'ayam') CLASS_AYAM = os.getenv("CLASS_AYAM", "ayam")
CLASS_TALENAN = os.getenv('CLASS_TALENAN', 'talenan') CLASS_TALENAN = os.getenv("CLASS_TALENAN", "talenan")
LINE_X = int(os.getenv('LINE_X')) if os.getenv('LINE_X') else None LINE_X = int(os.getenv("LINE_X")) if os.getenv("LINE_X") else None
LINE_X_FRAC = float(os.getenv('LINE_X_FRAC', '0.5')) LINE_X_FRAC = float(os.getenv("LINE_X_FRAC", "0.5"))
CROSS_DIRECTION = os.getenv('CROSS_DIRECTION', 'rtl').lower() CROSS_DIRECTION = os.getenv("CROSS_DIRECTION", "rtl").lower()
IMGSZ = int(os.getenv('IMGSZ', '320')) IMGSZ = int(os.getenv("IMGSZ", "320"))
HALF = os.getenv('HALF', 'false').lower() == 'true' HALF = os.getenv("HALF", "false").lower() == "true"
CONF = float(os.getenv('CONF', '0.3')) CONF = float(os.getenv("CONF", "0.3"))
DEVICE = int(os.getenv('DEVICE', '0')) DEVICE = int(os.getenv("DEVICE", "0"))
# RKNN NPU core mask # RKNN NPU core mask
CORE_MASK = int(os.getenv('CORE_MASK', '1')) CORE_MASK = int(os.getenv("CORE_MASK", "1"))
# YOLO decoder config # YOLO decoder config
NUM_CLASSES = int(os.getenv('NUM_CLASSES', '2')) NUM_CLASSES = int(os.getenv("NUM_CLASSES", "2"))
SCORE_SIGMOID = os.getenv('SCORE_SIGMOID', 'false').lower() == 'true' SCORE_SIGMOID = os.getenv("SCORE_SIGMOID", "false").lower() == "true"
# ByteTrack settings # ByteTrack settings
TRACK_HIGH_THRESH = float(os.getenv('TRACK_HIGH_THRESH', '0.5')) TRACK_HIGH_THRESH = float(os.getenv("TRACK_HIGH_THRESH", "0.5"))
TRACK_LOW_THRESH = float(os.getenv('TRACK_LOW_THRESH', '0.1')) TRACK_LOW_THRESH = float(os.getenv("TRACK_LOW_THRESH", "0.1"))
TRACK_MATCH_THRESH = float(os.getenv('TRACK_MATCH_THRESH', '0.8')) TRACK_MATCH_THRESH = float(os.getenv("TRACK_MATCH_THRESH", "0.8"))
TRACK_BUFFER = int(os.getenv('TRACK_BUFFER', '30')) TRACK_BUFFER = int(os.getenv("TRACK_BUFFER", "30"))
TRACK_MIN_HITS = int(os.getenv('TRACK_MIN_HITS', '3')) TRACK_MIN_HITS = int(os.getenv("TRACK_MIN_HITS", "3"))
DAILY_CUTOFF_TIME = os.getenv('DAILY_CUTOFF_TIME', '20:00') DAILY_CUTOFF_TIME = os.getenv("DAILY_CUTOFF_TIME", "20:00")
BATCH_TIMEOUT_SECONDS = float(os.getenv('BATCH_TIMEOUT_SECONDS', '300')) BATCH_TIMEOUT_SECONDS = float(os.getenv("BATCH_TIMEOUT_SECONDS", "300"))
IGNORE_BATCH_LABEL_TIMEOUT = float(os.getenv('IGNORE_BATCH_LABEL_TIMEOUT_SECONDS', '30')) IGNORE_BATCH_LABEL_TIMEOUT = float(
MIN_OBJECT_PER_BATCH = int(os.getenv('MIN_OBJECT_PER_BATCH', '60')) os.getenv("IGNORE_BATCH_LABEL_TIMEOUT_SECONDS", "30")
MIN_DURATION_PER_BATCH = int(os.getenv('MIN_DURATION_PER_BATCH', '60')) )
MIN_OBJECT_PER_BATCH = int(os.getenv("MIN_OBJECT_PER_BATCH", "60"))
MIN_DURATION_PER_BATCH = int(os.getenv("MIN_DURATION_PER_BATCH", "60"))
EXPORT_CSV = os.getenv('EXPORT_CSV', 'true').lower() == 'true' EXPORT_CSV = os.getenv("EXPORT_CSV", "true").lower() == "true"
CROSS_CSV = os.getenv('CROSS_CSV', f'{OUTPUT_DIR}/batch_crossings.csv') CROSS_CSV = os.getenv("CROSS_CSV", f"{OUTPUT_DIR}/batch_crossings.csv")
WARMUP_FRAMES = int(os.getenv('WARMUP_FRAMES', '30')) WARMUP_FRAMES = int(os.getenv("WARMUP_FRAMES", "30"))
RECONNECT_DELAY_SEC = int(os.getenv('RECONNECT_DELAY_SEC', '3')) RECONNECT_DELAY_SEC = int(os.getenv("RECONNECT_DELAY_SEC", "3"))
MAX_RECONNECT_ATTEMPTS = int(os.getenv('MAX_RECONNECT_ATTEMPTS', '0')) MAX_RECONNECT_ATTEMPTS = int(os.getenv("MAX_RECONNECT_ATTEMPTS", "0"))
FLUSH_EVERY_N_FRAMES = int(os.getenv('FLUSH_EVERY_N_FRAMES', '100')) FLUSH_EVERY_N_FRAMES = int(os.getenv("FLUSH_EVERY_N_FRAMES", "100"))
TRACKED_PRUNE_SEC = int(os.getenv('TRACKED_PRUNE_SEC', '300')) TRACKED_PRUNE_SEC = int(os.getenv("TRACKED_PRUNE_SEC", "300"))
RECORD_VIDEO = os.getenv('RECORD_VIDEO', 'false').lower() == 'true' RECORD_VIDEO = os.getenv("RECORD_VIDEO", "false").lower() == "true"
VIDEO_SEGMENT_SEC = int(os.getenv('VIDEO_SEGMENT_SEC', '3600')) VIDEO_SEGMENT_SEC = int(os.getenv("VIDEO_SEGMENT_SEC", "3600"))
OUTPUT_FPS = int(os.getenv('OUTPUT_FPS', '15')) OUTPUT_FPS = int(os.getenv("OUTPUT_FPS", "15"))
LIVE_STREAM_ENABLED = os.getenv('LIVE_STREAM_ENABLED', 'false').lower() == 'true' LIVE_STREAM_ENABLED = os.getenv("LIVE_STREAM_ENABLED", "false").lower() == "true"
LIVE_STREAM_FRAME_PATH = os.getenv('LIVE_STREAM_FRAME_PATH', '/dev/shm/jetson-counter/live_frame.jpg') LIVE_STREAM_FRAME_PATH = os.getenv(
LIVE_STREAM_QUALITY = int(os.getenv('LIVE_STREAM_QUALITY', '75')) "LIVE_STREAM_FRAME_PATH", "/dev/shm/jetson-counter/live_frame.jpg"
LIVE_STREAM_EVERY_N = int(os.getenv('LIVE_STREAM_EVERY_N', '2')) )
LIVE_STREAM_QUALITY = int(os.getenv("LIVE_STREAM_QUALITY", "75"))
LIVE_STREAM_EVERY_N = int(os.getenv("LIVE_STREAM_EVERY_N", "2"))
RTSP_FFMPEG_OPTIONS = os.getenv( RTSP_FFMPEG_OPTIONS = os.getenv(
'OPENCV_FFMPEG_CAPTURE_OPTIONS', "OPENCV_FFMPEG_CAPTURE_OPTIONS",
'rtsp_transport;tcp|fflags;nobuffer|flags;low_delay', "rtsp_transport;tcp|fflags;nobuffer|flags;low_delay",
) )
IS_LIVE = SOURCE.lower().startswith(('rtsp://', 'http://')) IS_LIVE = SOURCE.lower().startswith(("rtsp://", "http://"))
CROSS_FLASH_FRAMES = 12 CROSS_FLASH_FRAMES = 12
POPUP_LIFETIME = 20 POPUP_LIFETIME = 20
@@ -90,8 +96,14 @@ BATCH_PULSE_FRAMES = 20
SKELETON = [(0, 1), (4, 3), (1, 2), (3, 2), (2, 6), (2, 5), (2, 7), (7, 8)] SKELETON = [(0, 1), (4, 3), (1, 2), (3, 2), (2, 6), (2, 5), (2, 7), (7, 8)]
SK_COLORS = [ SK_COLORS = [
(0, 255, 255), (0, 255, 255), (255, 0, 255), (255, 0, 255), (0, 255, 255),
(0, 255, 0), (255, 255, 0), (0, 0, 255), (200, 200, 0), (0, 255, 255),
(255, 0, 255),
(255, 0, 255),
(0, 255, 0),
(255, 255, 0),
(0, 0, 255),
(200, 200, 0),
] ]
C_PANEL = (28, 24, 18) C_PANEL = (28, 24, 18)
@@ -111,7 +123,7 @@ shutdown_requested = False
def request_shutdown(signum, frame): def request_shutdown(signum, frame):
global shutdown_requested global shutdown_requested
shutdown_requested = True shutdown_requested = True
print('\nShutdown requested — finishing current frame...') print("\nShutdown requested — finishing current frame...")
signal.signal(signal.SIGINT, request_shutdown) signal.signal(signal.SIGINT, request_shutdown)
@@ -122,6 +134,7 @@ signal.signal(signal.SIGTERM, request_shutdown)
# YOLO output decoder (NMS only — boxes are pre-decoded by the model) # YOLO output decoder (NMS only — boxes are pre-decoded by the model)
# ============================================================================= # =============================================================================
def _nms(boxes, scores, iou_thr=0.45): def _nms(boxes, scores, iou_thr=0.45):
order = np.argsort(scores)[::-1] order = np.argsort(scores)[::-1]
keep = [] keep = []
@@ -138,7 +151,9 @@ def _nms(boxes, scores, iou_thr=0.45):
h = np.maximum(0.0, yy2 - yy1) h = np.maximum(0.0, yy2 - yy1)
inter = w * h inter = w * h
area_i = (boxes[idx, 2] - boxes[idx, 0]) * (boxes[idx, 3] - boxes[idx, 1]) area_i = (boxes[idx, 2] - boxes[idx, 0]) * (boxes[idx, 3] - boxes[idx, 1])
area_o = (boxes[order[1:], 2] - boxes[order[1:], 0]) * (boxes[order[1:], 3] - boxes[order[1:], 1]) area_o = (boxes[order[1:], 2] - boxes[order[1:], 0]) * (
boxes[order[1:], 3] - boxes[order[1:], 1]
)
iou = inter / (area_i + area_o - inter + 1e-16) iou = inter / (area_i + area_o - inter + 1e-16)
order = order[1:][iou < iou_thr] order = order[1:][iou < iou_thr]
return np.array(keep) return np.array(keep)
@@ -148,6 +163,7 @@ def _nms(boxes, scores, iou_thr=0.45):
# IoU helpers (xyxy format) # IoU helpers (xyxy format)
# ============================================================================= # =============================================================================
def _ious_xyxy(boxes_a, boxes_b): def _ious_xyxy(boxes_a, boxes_b):
"""Pairwise IoU: (N,4) vs (M,4) → (N,M) matrix.""" """Pairwise IoU: (N,4) vs (M,4) → (N,M) matrix."""
n, m = len(boxes_a), len(boxes_b) n, m = len(boxes_a), len(boxes_b)
@@ -190,6 +206,7 @@ def _greedy_match(cost_matrix, threshold=0.3):
# Kalman filter box tracker (state: x, y, w, h, vx, vy, vw, vh) # Kalman filter box tracker (state: x, y, w, h, vx, vy, vw, vh)
# ============================================================================= # =============================================================================
class KalmanBoxTracker: class KalmanBoxTracker:
count = 0 count = 0
@@ -242,6 +259,7 @@ class KalmanBoxTracker:
class _KalmanFilter: class _KalmanFilter:
"""8-state constant-velocity Kalman filter for bounding box tracking.""" """8-state constant-velocity Kalman filter for bounding box tracking."""
def __init__(self): def __init__(self):
ndim, dt = 4, 1.0 ndim, dt = 4, 1.0
self.motion_mat = np.eye(2 * ndim, 2 * ndim, dtype=np.float32) self.motion_mat = np.eye(2 * ndim, 2 * ndim, dtype=np.float32)
@@ -274,12 +292,16 @@ class _KalmanFilter:
self.P = self.motion_mat @ self.P @ self.motion_mat.T + Q self.P = self.motion_mat @ self.P @ self.motion_mat.T + Q
def update(self, z): def update(self, z):
R = np.diag(np.square([ R = np.diag(
self._std_weight_position * z[2], np.square(
self._std_weight_position * z[3], [
self._std_weight_position * z[2], self._std_weight_position * z[2],
self._std_weight_position * z[3], self._std_weight_position * z[3],
])) self._std_weight_position * z[2],
self._std_weight_position * z[3],
]
)
)
H = self.update_mat H = self.update_mat
S = H @ self.P @ H.T + R S = H @ self.P @ H.T + R
@@ -295,11 +317,18 @@ class _KalmanFilter:
# ByteTrack multi-object tracker # ByteTrack multi-object tracker
# ============================================================================= # =============================================================================
class ByteTracker: class ByteTracker:
"""ByteTrack: two-stage association with Kalman filter prediction.""" """ByteTrack: two-stage association with Kalman filter prediction."""
def __init__(self, track_high_thresh=0.5, track_low_thresh=0.1, def __init__(
match_thresh=0.8, track_buffer=30, min_hits=3): self,
track_high_thresh=0.5,
track_low_thresh=0.1,
match_thresh=0.8,
track_buffer=30,
min_hits=3,
):
self.high_thresh = track_high_thresh self.high_thresh = track_high_thresh
self.low_thresh = track_low_thresh self.low_thresh = track_low_thresh
self.match_thresh = match_thresh self.match_thresh = match_thresh
@@ -367,7 +396,9 @@ class ByteTracker:
tracked_map[track_pool[ti].track_id] = track_pool[ti].get_cx() tracked_map[track_pool[ti].track_id] = track_pool[ti].get_cx()
match_pairs_high.append((det_global, ti)) match_pairs_high.append((det_global, ti))
unmatched_tracks = [t for t in range(num_tracks) if t not in matched_track_idx] unmatched_tracks = [
t for t in range(num_tracks) if t not in matched_track_idx
]
# --- second association: low-score ↔ unmatched tracks --- # --- second association: low-score ↔ unmatched tracks ---
low_idx = np.where(is_low)[0] low_idx = np.where(is_low)[0]
@@ -383,10 +414,14 @@ class ByteTracker:
det_global = int(low_idx[di]) det_global = int(low_idx[di])
pool_idx = unmatched_tracks[uti] pool_idx = unmatched_tracks[uti]
track_pool[pool_idx].update(dets[det_global]) track_pool[pool_idx].update(dets[det_global])
track_pool[pool_idx].hit_streak = max(1, track_pool[pool_idx].hit_streak) track_pool[pool_idx].hit_streak = max(
1, track_pool[pool_idx].hit_streak
)
matched_track_idx.add(pool_idx) matched_track_idx.add(pool_idx)
det_to_track[det_global] = track_pool[pool_idx].track_id det_to_track[det_global] = track_pool[pool_idx].track_id
tracked_map[track_pool[pool_idx].track_id] = track_pool[pool_idx].get_cx() tracked_map[track_pool[pool_idx].track_id] = track_pool[
pool_idx
].get_cx()
# --- reset hit_streak for unmatched tracks --- # --- reset hit_streak for unmatched tracks ---
for ti, trk in enumerate(track_pool): for ti, trk in enumerate(track_pool):
@@ -432,9 +467,19 @@ class ByteTracker:
# RKNN YOLO wrapper (detect output format: (1, 4+num_classes, N)) # RKNN YOLO wrapper (detect output format: (1, 4+num_classes, N))
# ============================================================================= # =============================================================================
class RKNNYOLO: class RKNNYOLO:
def __init__(self, model_path, core_mask=1, imgsz=320, conf=0.3, iou=0.45, def __init__(
num_classes=2, num_keypoints=0, score_sigmoid=False): self,
model_path,
core_mask=1,
imgsz=320,
conf=0.3,
iou=0.45,
num_classes=2,
num_keypoints=0,
score_sigmoid=False,
):
self.imgsz = imgsz self.imgsz = imgsz
self.conf = conf self.conf = conf
self.iou = iou self.iou = iou
@@ -445,18 +490,18 @@ class RKNNYOLO:
self.rknn = RKNNLite(verbose=False) self.rknn = RKNNLite(verbose=False)
ret = self.rknn.load_rknn(model_path) ret = self.rknn.load_rknn(model_path)
if ret != 0: if ret != 0:
raise RuntimeError(f'Failed to load RKNN model: {model_path}') raise RuntimeError(f"Failed to load RKNN model: {model_path}")
ret = self.rknn.init_runtime(core_mask=core_mask) ret = self.rknn.init_runtime(core_mask=core_mask)
if ret != 0: if ret != 0:
raise RuntimeError(f'Failed to init RKNN runtime (core_mask={core_mask})') raise RuntimeError(f"Failed to init RKNN runtime (core_mask={core_mask})")
try: try:
sdk_ver = self.rknn.get_sdk_version() sdk_ver = self.rknn.get_sdk_version()
print(f'RKNN SDK version: {sdk_ver}') print(f"RKNN SDK version: {sdk_ver}")
except Exception: except Exception:
pass pass
print(f'RKNN model loaded: {model_path} imgsz={imgsz} core_mask={core_mask}') print(f"RKNN model loaded: {model_path} imgsz={imgsz} core_mask={core_mask}")
def _preprocess(self, frame): def _preprocess(self, frame):
h0, w0 = frame.shape[:2] h0, w0 = frame.shape[:2]
@@ -467,7 +512,7 @@ class RKNNYOLO:
letterbox = np.full((self.imgsz, self.imgsz, 3), 114, dtype=np.uint8) letterbox = np.full((self.imgsz, self.imgsz, 3), 114, dtype=np.uint8)
dy = (self.imgsz - nh) // 2 dy = (self.imgsz - nh) // 2
dx = (self.imgsz - nw) // 2 dx = (self.imgsz - nw) // 2
letterbox[dy:dy + nh, dx:dx + nw] = resized letterbox[dy : dy + nh, dx : dx + nw] = resized
rgb = cv2.cvtColor(letterbox, cv2.COLOR_BGR2RGB) rgb = cv2.cvtColor(letterbox, cv2.COLOR_BGR2RGB)
gains = np.array([scale, scale, dy, dx], dtype=np.float32) gains = np.array([scale, scale, dy, dx], dtype=np.float32)
@@ -500,12 +545,15 @@ class RKNNYOLO:
else: else:
cls_scores = cls_raw cls_scores = cls_raw
boxes_xyxy = np.stack([ boxes_xyxy = np.stack(
boxes_cxcywh[:, 0] - boxes_cxcywh[:, 2] / 2, [
boxes_cxcywh[:, 1] - boxes_cxcywh[:, 3] / 2, boxes_cxcywh[:, 0] - boxes_cxcywh[:, 2] / 2,
boxes_cxcywh[:, 0] + boxes_cxcywh[:, 2] / 2, boxes_cxcywh[:, 1] - boxes_cxcywh[:, 3] / 2,
boxes_cxcywh[:, 1] + boxes_cxcywh[:, 3] / 2, boxes_cxcywh[:, 0] + boxes_cxcywh[:, 2] / 2,
], axis=1) boxes_cxcywh[:, 1] + boxes_cxcywh[:, 3] / 2,
],
axis=1,
)
max_scores = cls_scores.max(axis=1) max_scores = cls_scores.max(axis=1)
class_ids = cls_scores.argmax(axis=1) class_ids = cls_scores.argmax(axis=1)
@@ -536,12 +584,14 @@ class RKNNYOLO:
keep = _nms(bboxes[idx], scores[idx], iou_thr=self.iou) keep = _nms(bboxes[idx], scores[idx], iou_thr=self.iou)
for k in keep: for k in keep:
j = idx[k] j = idx[k]
detections.append({ detections.append(
'bbox': bboxes[j].tolist(), {
'score': float(scores[j]), "bbox": bboxes[j].tolist(),
'cls': int(clses[j]), "score": float(scores[j]),
'keypoints': None, "cls": int(clses[j]),
}) "keypoints": None,
}
)
return detections return detections
def release(self): def release(self):
@@ -552,6 +602,7 @@ class RKNNYOLO:
# Drawing helpers # Drawing helpers
# ============================================================================= # =============================================================================
def resolve_line_x(frame_width): def resolve_line_x(frame_width):
if LINE_X is not None: if LINE_X is not None:
return LINE_X return LINE_X
@@ -561,41 +612,41 @@ def resolve_line_x(frame_width):
def crossed_line(prev_cx, cx, line_x, direction=CROSS_DIRECTION): def crossed_line(prev_cx, cx, line_x, direction=CROSS_DIRECTION):
if direction == 'ltr': if direction == "ltr":
return prev_cx < line_x <= cx return prev_cx < line_x <= cx
if direction == 'both': if direction == "both":
return (prev_cx > line_x >= cx) or (prev_cx < line_x <= cx) return (prev_cx > line_x >= cx) or (prev_cx < line_x <= cx)
return prev_cx > line_x >= cx return prev_cx > line_x >= cx
def now_str(): def now_str():
return datetime.now().strftime('%Y-%m-%d %H:%M:%S') return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
def open_capture(source): def open_capture(source):
if source.lower().startswith(('rtsp://', 'http://')): if source.lower().startswith(("rtsp://", "http://")):
os.environ['OPENCV_FFMPEG_CAPTURE_OPTIONS'] = RTSP_FFMPEG_OPTIONS os.environ["OPENCV_FFMPEG_CAPTURE_OPTIONS"] = RTSP_FFMPEG_OPTIONS
cap = cv2.VideoCapture(source, cv2.CAP_FFMPEG) cap = cv2.VideoCapture(source, cv2.CAP_FFMPEG)
cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
return cap return cap
def warmup_stream(cap, n=WARMUP_FRAMES): def warmup_stream(cap, n=WARMUP_FRAMES):
print('Warming up stream...') print("Warming up stream...")
for _ in range(n): for _ in range(n):
cap.read() cap.read()
print('Stream ready!') print("Stream ready!")
def open_video_writer(path, w, h, fps): def open_video_writer(path, w, h, fps):
return cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*'avc1'), fps, (w, h)) return cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*"avc1"), fps, (w, h))
class CsvLogger: class CsvLogger:
def __init__(self, path, header): def __init__(self, path, header):
Path(path).parent.mkdir(parents=True, exist_ok=True) Path(path).parent.mkdir(parents=True, exist_ok=True)
new_file = not Path(path).exists() or Path(path).stat().st_size == 0 new_file = not Path(path).exists() or Path(path).stat().st_size == 0
self.file = open(path, 'a', newline='', buffering=1) self.file = open(path, "a", newline="", buffering=1)
self.writer = csv.writer(self.file) self.writer = csv.writer(self.file)
if new_file: if new_file:
self.writer.writerow(header) self.writer.writerow(header)
@@ -620,8 +671,8 @@ class VideoSegmentWriter:
self._open_next() self._open_next()
def _segment_path(self): def _segment_path(self):
ts = datetime.now().strftime('%Y%m%d_%H%M%S') ts = datetime.now().strftime("%Y%m%d_%H%M%S")
return str(self.output_dir / f'live_{ts}.mp4') return str(self.output_dir / f"live_{ts}.mp4")
def _open_next(self): def _open_next(self):
if self.writer is not None: if self.writer is not None:
@@ -629,7 +680,7 @@ class VideoSegmentWriter:
path = self._segment_path() path = self._segment_path()
self.writer = open_video_writer(path, self.w, self.h, self.fps) self.writer = open_video_writer(path, self.w, self.h, self.fps)
self.segment_start = time.monotonic() self.segment_start = time.monotonic()
print(f'Recording segment: {path}') print(f"Recording segment: {path}")
def write(self, frame): def write(self, frame):
if time.monotonic() - self.segment_start >= self.segment_sec: if time.monotonic() - self.segment_start >= self.segment_sec:
@@ -642,7 +693,9 @@ class VideoSegmentWriter:
def prune_stale_tracks(tracked, now_mono): def prune_stale_tracks(tracked, now_mono):
stale = [tid for tid, (_, ts) in tracked.items() if now_mono - ts > TRACKED_PRUNE_SEC] stale = [
tid for tid, (_, ts) in tracked.items() if now_mono - ts > TRACKED_PRUNE_SEC
]
for tid in stale: for tid in stale:
del tracked[tid] del tracked[tid]
@@ -680,7 +733,16 @@ def draw_elegant_counting_line(img, line_x, h, pulse_remaining=0):
y_end = min(y + dash_len, h) y_end = min(y + dash_len, h)
cv2.line(img, (line_x, y), (line_x, y_end), C_LINE_CORE, 2, cv2.LINE_AA) cv2.line(img, (line_x, y), (line_x, y_end), C_LINE_CORE, 2, cv2.LINE_AA)
y += dash_len + gap y += dash_len + gap
cv2.putText(img, 'COUNT LINE', (line_x - 46, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.42, C_LINE_CORE, 1, cv2.LINE_AA) cv2.putText(
img,
"COUNT LINE",
(line_x - 46, 24),
cv2.FONT_HERSHEY_SIMPLEX,
0.42,
C_LINE_CORE,
1,
cv2.LINE_AA,
)
def draw_hero_count(img, line_x, h, count, pulse_remaining=0): def draw_hero_count(img, line_x, h, count, pulse_remaining=0):
@@ -691,34 +753,120 @@ def draw_hero_count(img, line_x, h, count, pulse_remaining=0):
(tw, th), _ = cv2.getTextSize(text, font, font_scale, thickness) (tw, th), _ = cv2.getTextSize(text, font, font_scale, thickness)
pad = 14 pad = 14
tx, ty = line_x - tw // 2, h // 2 + th // 2 tx, ty = line_x - tw // 2, h // 2 + th // 2
overlay_rect(img, tx - pad, ty - th - pad, tx + tw + pad, ty + pad // 2, C_PANEL, alpha=0.78) overlay_rect(
cv2.rectangle(img, (tx - pad, ty - th - pad), (tx + tw + pad, ty + pad // 2), C_LINE_CORE, 2) img, tx - pad, ty - th - pad, tx + tw + pad, ty + pad // 2, C_PANEL, alpha=0.78
)
cv2.rectangle(
img, (tx - pad, ty - th - pad), (tx + tw + pad, ty + pad // 2), C_LINE_CORE, 2
)
cv2.putText(img, text, (tx, ty), font, font_scale, C_GREEN, thickness, cv2.LINE_AA) cv2.putText(img, text, (tx, ty), font, font_scale, C_GREEN, thickness, cv2.LINE_AA)
def draw_hud(img, w, batch_num, batch_count, total_ayam, elapsed_sec, rate, camera_id, clock): def draw_hud(
img, w, batch_num, batch_count, total_ayam, elapsed_sec, rate, camera_id, clock
):
bar_h = 52 bar_h = 52
overlay_rect(img, 0, 0, w, bar_h, C_PANEL, alpha=0.72) overlay_rect(img, 0, 0, w, bar_h, C_PANEL, alpha=0.72)
cv2.line(img, (0, bar_h), (w, bar_h), C_BORDER, 1) cv2.line(img, (0, bar_h), (w, bar_h), C_BORDER, 1)
cv2.putText(img, 'BATCH', (16, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) cv2.putText(
batch_label = str(batch_num) if batch_num else '\u2014' img, "BATCH", (16, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA
cv2.putText(img, batch_label, (16, 44), cv2.FONT_HERSHEY_SIMPLEX, 0.9, C_ACCENT, 2, cv2.LINE_AA) )
cv2.putText(img, 'COUNT', (100, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) batch_label = str(batch_num) if batch_num else "\u2014"
cv2.putText(img, str(batch_count), (100, 44), cv2.FONT_HERSHEY_SIMPLEX, 0.9, C_GREEN, 2, cv2.LINE_AA) cv2.putText(
cv2.putText(img, 'TOTAL', (190, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) img,
cv2.putText(img, str(total_ayam), (190, 44), cv2.FONT_HERSHEY_SIMPLEX, 0.7, C_TEXT, 1, cv2.LINE_AA) batch_label,
cv2.putText(img, 'UPTIME', (280, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) (16, 44),
cv2.putText(img, f'{elapsed_sec / 3600:.1f}h', (280, 44), cv2.FONT_HERSHEY_SIMPLEX, 0.7, C_TEXT, 1, cv2.LINE_AA) cv2.FONT_HERSHEY_SIMPLEX,
cv2.putText(img, 'RATE', (380, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) 0.9,
cv2.putText(img, f'{rate:.1f}/min', (380, 44), cv2.FONT_HERSHEY_SIMPLEX, 0.7, C_ACCENT, 1, cv2.LINE_AA) C_ACCENT,
cv2.putText(img, clock, (w - 180, 36), cv2.FONT_HERSHEY_SIMPLEX, 0.55, C_TEXT, 1, cv2.LINE_AA) 2,
cv2.putText(img, f'CAM {camera_id}', (w - 180, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) cv2.LINE_AA,
)
cv2.putText(
img, "COUNT", (100, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA
)
cv2.putText(
img,
str(batch_count),
(100, 44),
cv2.FONT_HERSHEY_SIMPLEX,
0.9,
C_GREEN,
2,
cv2.LINE_AA,
)
cv2.putText(
img, "TOTAL", (190, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA
)
cv2.putText(
img,
str(total_ayam),
(190, 44),
cv2.FONT_HERSHEY_SIMPLEX,
0.7,
C_TEXT,
1,
cv2.LINE_AA,
)
cv2.putText(
img,
"UPTIME",
(280, 20),
cv2.FONT_HERSHEY_SIMPLEX,
0.45,
C_MUTED,
1,
cv2.LINE_AA,
)
cv2.putText(
img,
f"{elapsed_sec / 3600:.1f}h",
(280, 44),
cv2.FONT_HERSHEY_SIMPLEX,
0.7,
C_TEXT,
1,
cv2.LINE_AA,
)
cv2.putText(
img, "RATE", (380, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA
)
cv2.putText(
img,
f"{rate:.1f}/min",
(380, 44),
cv2.FONT_HERSHEY_SIMPLEX,
0.7,
C_ACCENT,
1,
cv2.LINE_AA,
)
# cv2.putText(img, clock, (w - 180, 36), cv2.FONT_HERSHEY_SIMPLEX, 0.55, C_TEXT, 1, cv2.LINE_AA)
cv2.putText(
img,
f"CAM {camera_id}",
(w - 180, 20),
cv2.FONT_HERSHEY_SIMPLEX,
0.45,
C_MUTED,
1,
cv2.LINE_AA,
)
def draw_footer(img, w, h, frame_idx, live_tag): def draw_footer(img, w, h, frame_idx, live_tag):
bar_h = 28 bar_h = 28
overlay_rect(img, 0, h - bar_h, w, h, C_PANEL, alpha=0.55) overlay_rect(img, 0, h - bar_h, w, h, C_PANEL, alpha=0.55)
cv2.putText(img, f'{live_tag} | Frame {frame_idx}', (12, h - 9), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) cv2.putText(
img,
f"{live_tag} | Frame {frame_idx}",
(12, h - 9),
cv2.FONT_HERSHEY_SIMPLEX,
0.45,
C_MUTED,
1,
cv2.LINE_AA,
)
def draw_skeleton_bold(img, kpts): def draw_skeleton_bold(img, kpts):
@@ -738,28 +886,39 @@ def draw_skeleton_bold(img, kpts):
def draw_popups(img, popups, frame_idx): def draw_popups(img, popups, frame_idx):
alive = [] alive = []
for pop in popups: for pop in popups:
age = frame_idx - pop['born'] age = frame_idx - pop["born"]
if age > POPUP_LIFETIME: if age > POPUP_LIFETIME:
continue continue
alive.append(pop) alive.append(pop)
fade = 1.0 - age / POPUP_LIFETIME fade = 1.0 - age / POPUP_LIFETIME
y = pop['y'] - int(age * 1.8) y = pop["y"] - int(age * 1.8)
color = (int(C_GREEN[0] * fade), int(C_GREEN[1] * fade), int(C_GREEN[2] * fade)) color = (int(C_GREEN[0] * fade), int(C_GREEN[1] * fade), int(C_GREEN[2] * fade))
cv2.putText(img, pop['text'], (pop['x'], y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2, cv2.LINE_AA) cv2.putText(
img,
pop["text"],
(pop["x"], y),
cv2.FONT_HERSHEY_SIMPLEX,
0.7,
color,
2,
cv2.LINE_AA,
)
return alive return alive
def draw_batch_banner(img, w, batch_num, pulse_remaining): def draw_batch_banner(img, w, batch_num, pulse_remaining):
if pulse_remaining <= 0: if pulse_remaining <= 0:
return return
text = f'NEW BATCH {batch_num}' text = f"NEW BATCH {batch_num}"
font = cv2.FONT_HERSHEY_SIMPLEX font = cv2.FONT_HERSHEY_SIMPLEX
(tw, th), _ = cv2.getTextSize(text, font, 0.8, 2) (tw, th), _ = cv2.getTextSize(text, font, 0.8, 2)
x1, y1 = w // 2 - tw // 2 - 16, 62 x1, y1 = w // 2 - tw // 2 - 16, 62
x2, y2 = w // 2 + tw // 2 + 16, 62 + th + 20 x2, y2 = w // 2 + tw // 2 + 16, 62 + th + 20
overlay_rect(img, x1, y1, x2, y2, C_PANEL, alpha=0.7) overlay_rect(img, x1, y1, x2, y2, C_PANEL, alpha=0.7)
cv2.rectangle(img, (x1, y1), (x2, y2), C_ACCENT, 2) cv2.rectangle(img, (x1, y1), (x2, y2), C_ACCENT, 2)
cv2.putText(img, text, (w // 2 - tw // 2, 62 + th + 4), font, 0.8, C_ACCENT, 2, cv2.LINE_AA) cv2.putText(
img, text, (w // 2 - tw // 2, 62 + th + 4), font, 0.8, C_ACCENT, 2, cv2.LINE_AA
)
def connect_stream(source, warmup=WARMUP_FRAMES): def connect_stream(source, warmup=WARMUP_FRAMES):
@@ -769,11 +928,13 @@ def connect_stream(source, warmup=WARMUP_FRAMES):
if not cap.isOpened(): if not cap.isOpened():
attempts += 1 attempts += 1
if MAX_RECONNECT_ATTEMPTS and attempts >= MAX_RECONNECT_ATTEMPTS: if MAX_RECONNECT_ATTEMPTS and attempts >= MAX_RECONNECT_ATTEMPTS:
raise RuntimeError(f'Cannot open source after {attempts} attempts: {source}') raise RuntimeError(
print(f'Cannot open source, retry in {RECONNECT_DELAY_SEC}s...') f"Cannot open source after {attempts} attempts: {source}"
)
print(f"Cannot open source, retry in {RECONNECT_DELAY_SEC}s...")
time.sleep(RECONNECT_DELAY_SEC) time.sleep(RECONNECT_DELAY_SEC)
continue continue
if warmup > 0 and source.lower().startswith(('rtsp://', 'http://')): if warmup > 0 and source.lower().startswith(("rtsp://", "http://")):
warmup_stream(cap, warmup) warmup_stream(cap, warmup)
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
@@ -788,6 +949,7 @@ def connect_stream(source, warmup=WARMUP_FRAMES):
# Main loop # Main loop
# ============================================================================= # =============================================================================
def run(): def run():
global shutdown_requested global shutdown_requested
@@ -801,13 +963,15 @@ def run():
ignore_batch_label_timeout=IGNORE_BATCH_LABEL_TIMEOUT, ignore_batch_label_timeout=IGNORE_BATCH_LABEL_TIMEOUT,
min_object_per_batch=MIN_OBJECT_PER_BATCH, min_object_per_batch=MIN_OBJECT_PER_BATCH,
min_duration_per_batch=MIN_DURATION_PER_BATCH, min_duration_per_batch=MIN_DURATION_PER_BATCH,
logger=lambda msg: print(f'[{now_str()}] {msg}'), logger=lambda msg: print(f"[{now_str()}] {msg}"),
) )
store.start_cutoff_watcher() store.start_cutoff_watcher()
cross_logger = None cross_logger = None
if EXPORT_CSV: if EXPORT_CSV:
cross_logger = CsvLogger(CROSS_CSV, ['batch', 'frame', 'timestamp', 'chicken_id']) cross_logger = CsvLogger(
CROSS_CSV, ["batch", "frame", "timestamp", "chicken_id"]
)
model = RKNNYOLO( model = RKNNYOLO(
model_path=MODEL_PATH, model_path=MODEL_PATH,
@@ -819,8 +983,8 @@ def run():
) )
CLASS_IDS = { CLASS_IDS = {
os.getenv('CLASS_AYAM', 'ayam'): 0, os.getenv("CLASS_AYAM", "ayam"): 0,
os.getenv('CLASS_TALENAN', 'talenan'): 1, os.getenv("CLASS_TALENAN", "talenan"): 1,
} }
ayam_cls = CLASS_IDS[CLASS_AYAM] ayam_cls = CLASS_IDS[CLASS_AYAM]
talenan_cls = CLASS_IDS[CLASS_TALENAN] talenan_cls = CLASS_IDS[CLASS_TALENAN]
@@ -859,12 +1023,16 @@ def run():
return return
line_x = resolve_line_x(w) line_x = resolve_line_x(w)
print(f'RKNN+ByteTrack counter | {w}x{h} @ {fps}fps | line x={line_x} | cross={CROSS_DIRECTION}') print(
print(f'Model: {MODEL_PATH} | imgsz={IMGSZ} | core_mask={CORE_MASK}') f"RKNN+ByteTrack counter | {w}x{h} @ {fps}fps | line x={line_x} | cross={CROSS_DIRECTION}"
print(f'ByteTrack: high_thresh={TRACK_HIGH_THRESH} low_thresh={TRACK_LOW_THRESH} ' )
f'match_thresh={TRACK_MATCH_THRESH} buffer={TRACK_BUFFER}') print(f"Model: {MODEL_PATH} | imgsz={IMGSZ} | core_mask={CORE_MASK}")
print(f'DB: {DB_PATH}') print(
print(f'State: {STATE_FILE}') f"ByteTrack: high_thresh={TRACK_HIGH_THRESH} low_thresh={TRACK_LOW_THRESH} "
f"match_thresh={TRACK_MATCH_THRESH} buffer={TRACK_BUFFER}"
)
print(f"DB: {DB_PATH}")
print(f"State: {STATE_FILE}")
if RECORD_VIDEO: if RECORD_VIDEO:
video_writer = VideoSegmentWriter(OUTPUT_DIR, w, h, fps, VIDEO_SEGMENT_SEC) video_writer = VideoSegmentWriter(OUTPUT_DIR, w, h, fps, VIDEO_SEGMENT_SEC)
@@ -877,7 +1045,9 @@ def run():
if not IS_LIVE: if not IS_LIVE:
break break
reconnect_count += 1 reconnect_count += 1
print(f'Stream dropped (attempt {reconnect_count}), reconnecting in {RECONNECT_DELAY_SEC}s...') print(
f"Stream dropped (attempt {reconnect_count}), reconnecting in {RECONNECT_DELAY_SEC}s..."
)
cap.release() cap.release()
time.sleep(RECONNECT_DELAY_SEC) time.sleep(RECONNECT_DELAY_SEC)
cap, w, h, fps = connect_stream(SOURCE) cap, w, h, fps = connect_stream(SOURCE)
@@ -905,10 +1075,10 @@ def run():
talenan_cx_list = [] talenan_cx_list = []
for det in detections: for det in detections:
bbox = det['bbox'] bbox = det["bbox"]
score = det['score'] score = det["score"]
cls_id = det['cls'] cls_id = det["cls"]
kpts = det['keypoints'] kpts = det["keypoints"]
cx = (bbox[0] + bbox[2]) / 2.0 cx = (bbox[0] + bbox[2]) / 2.0
if cls_id == talenan_cls: if cls_id == talenan_cls:
@@ -928,9 +1098,11 @@ def run():
talenan_scores = np.array(talenan_scores, dtype=np.float32) talenan_scores = np.array(talenan_scores, dtype=np.float32)
ayam_track_map, ayam_det_to_track, ayam_lost_map = ayam_tracker.update( ayam_track_map, ayam_det_to_track, ayam_lost_map = ayam_tracker.update(
ayam_boxes_xyxy, ayam_scores) ayam_boxes_xyxy, ayam_scores
talenan_track_map, talenan_det_to_track, talenan_lost_map = talenan_tracker.update( )
talenan_boxes_xyxy, talenan_scores) talenan_track_map, talenan_det_to_track, talenan_lost_map = (
talenan_tracker.update(talenan_boxes_xyxy, talenan_scores)
)
# Process talenan crossings # Process talenan crossings
for di in range(len(talenan_boxes_xyxy)): for di in range(len(talenan_boxes_xyxy)):
@@ -942,17 +1114,22 @@ def run():
if tid in talenan_tracked: if tid in talenan_tracked:
prev_cx = talenan_tracked[tid][0] prev_cx = talenan_tracked[tid][0]
if crossed_line(prev_cx, cx, line_x) and tid not in talenan_line_crossed: if (
crossed_line(prev_cx, cx, line_x)
and tid not in talenan_line_crossed
):
talenan_line_crossed.add(tid) talenan_line_crossed.add(tid)
if store.record_talenan_crossing(tid): if store.record_talenan_crossing(tid):
batch_closed_frame = True batch_closed_frame = True
talenan_cross_flash[tid] = CROSS_FLASH_FRAMES talenan_cross_flash[tid] = CROSS_FLASH_FRAMES
popups.append({ popups.append(
'x': int(cx) - 20, {
'y': int((bbox[1] + bbox[3]) / 2), "x": int(cx) - 20,
'born': frame_idx, "y": int((bbox[1] + bbox[3]) / 2),
'text': 'BATCH CLOSED', "born": frame_idx,
}) "text": "BATCH CLOSED",
}
)
talenan_tracked[tid] = (cx, mono) talenan_tracked[tid] = (cx, mono)
# Process ayam crossings (including lost tracks for line-cross continuity) # Process ayam crossings (including lost tracks for line-cross continuity)
@@ -964,24 +1141,36 @@ def run():
if tid in ayam_tracked: if tid in ayam_tracked:
prev_cx = ayam_tracked[tid][0] prev_cx = ayam_tracked[tid][0]
if crossed_line(prev_cx, cx, line_x) and tid not in ayam_line_crossed: if (
crossed_line(prev_cx, cx, line_x)
and tid not in ayam_line_crossed
):
ayam_line_crossed.add(tid) ayam_line_crossed.add(tid)
_, started_new = store.record_ayam_crossing(tid) _, started_new = store.record_ayam_crossing(tid)
if cross_logger: if cross_logger:
cross_logger.write_row([ cross_logger.write_row(
store.current_batch_number, frame_idx, [
datetime.now().isoformat(), tid, store.current_batch_number,
]) frame_idx,
datetime.now().isoformat(),
tid,
]
)
ayam_crossed_frame = True ayam_crossed_frame = True
if started_new: if started_new:
batch_started_frame = True batch_started_frame = True
ayam_cross_flash[tid] = CROSS_FLASH_FRAMES ayam_cross_flash[tid] = CROSS_FLASH_FRAMES
popups.append({ popups.append(
'x': int(cx) - 12, {
'y': int((ayam_boxes_xyxy[di][1] + ayam_boxes_xyxy[di][3]) / 2), "x": int(cx) - 12,
'born': frame_idx, "y": int(
'text': '+1', (ayam_boxes_xyxy[di][1] + ayam_boxes_xyxy[di][3])
}) / 2
),
"born": frame_idx,
"text": "+1",
}
)
ayam_tracked[tid] = (cx, mono) ayam_tracked[tid] = (cx, mono)
# Also track lost tracks for line-crossing continuity # Also track lost tracks for line-crossing continuity
@@ -999,7 +1188,7 @@ def run():
flash = talenan_cross_flash.get(tid, 0) flash = talenan_cross_flash.get(tid, 0)
color = C_GREEN if flash > 0 else C_TALENAN_BOX color = C_GREEN if flash > 0 else C_TALENAN_BOX
cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3 if flash > 0 else 2) cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3 if flash > 0 else 2)
draw_pill(frame, f'TALENAN {tid}', x1, y1 - 4, color) draw_pill(frame, f"TALENAN {tid}", x1, y1 - 4, color)
# Draw ayam # Draw ayam
for di in range(len(ayam_boxes_xyxy)): for di in range(len(ayam_boxes_xyxy)):
@@ -1011,7 +1200,7 @@ def run():
flash = ayam_cross_flash.get(tid, 0) flash = ayam_cross_flash.get(tid, 0)
color = C_GREEN if flash > 0 else C_AYAM_BOX color = C_GREEN if flash > 0 else C_AYAM_BOX
cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3 if flash > 0 else 2) cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3 if flash > 0 else 2)
draw_pill(frame, f'ID {tid}', x1, y1 - 4, color) draw_pill(frame, f"ID {tid}", x1, y1 - 4, color)
kpts = ayam_kpts_list[di] if di < len(ayam_kpts_list) else None kpts = ayam_kpts_list[di] if di < len(ayam_kpts_list) else None
if kpts is not None: if kpts is not None:
draw_skeleton_bold(frame, kpts) draw_skeleton_bold(frame, kpts)
@@ -1031,9 +1220,21 @@ def run():
draw_elegant_counting_line(frame, line_x, h, line_pulse) draw_elegant_counting_line(frame, line_x, h, line_pulse)
draw_hero_count(frame, line_x, h, batch_count, count_pulse) draw_hero_count(frame, line_x, h, batch_count, count_pulse)
draw_hud(frame, w, batch_num, batch_count, display_total, elapsed, rate, CAMERA_NAME, now_str()) draw_hud(
frame,
w,
batch_num,
batch_count,
display_total,
elapsed,
rate,
CAMERA_NAME,
now_str(),
)
draw_batch_banner(frame, w, batch_num, batch_pulse) draw_batch_banner(frame, w, batch_num, batch_pulse)
draw_footer(frame, w, h, frame_idx, 'LIVE-RKNN-BT' if IS_LIVE else 'FILE-RKNN-BT') draw_footer(
frame, w, h, frame_idx, "LIVE-RKNN-BT" if IS_LIVE else "FILE-RKNN-BT"
)
popups = draw_popups(frame, popups, frame_idx) popups = draw_popups(frame, popups, frame_idx)
for flash_store in (ayam_cross_flash, talenan_cross_flash): for flash_store in (ayam_cross_flash, talenan_cross_flash):
@@ -1051,8 +1252,10 @@ def run():
if LIVE_STREAM_ENABLED and frame_idx % LIVE_STREAM_EVERY_N == 0: if LIVE_STREAM_ENABLED and frame_idx % LIVE_STREAM_EVERY_N == 0:
try: try:
Path(LIVE_STREAM_FRAME_PATH).parent.mkdir(parents=True, exist_ok=True) Path(LIVE_STREAM_FRAME_PATH).parent.mkdir(parents=True, exist_ok=True)
_, jpeg = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, LIVE_STREAM_QUALITY]) _, jpeg = cv2.imencode(
with open(LIVE_STREAM_FRAME_PATH, 'wb') as f: ".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, LIVE_STREAM_QUALITY]
)
with open(LIVE_STREAM_FRAME_PATH, "wb") as f:
f.write(jpeg.tobytes()) f.write(jpeg.tobytes())
except Exception: except Exception:
pass pass
@@ -1060,8 +1263,8 @@ def run():
frame_idx += 1 frame_idx += 1
if frame_idx % FLUSH_EVERY_N_FRAMES == 0: if frame_idx % FLUSH_EVERY_N_FRAMES == 0:
print( print(
f'[{now_str()}] Frame {frame_idx} | Batch {batch_num}: {batch_count} ' f"[{now_str()}] Frame {frame_idx} | Batch {batch_num}: {batch_count} "
f'| Total: {display_total} | Uptime {elapsed / 3600:.2f}h' f"| Total: {display_total} | Uptime {elapsed / 3600:.2f}h"
) )
prune_stale_tracks(ayam_tracked, mono) prune_stale_tracks(ayam_tracked, mono)
prune_stale_tracks(talenan_tracked, mono) prune_stale_tracks(talenan_tracked, mono)
@@ -1074,13 +1277,13 @@ def run():
model.release() model.release()
store.shutdown() store.shutdown()
print('\n=== Batch Summary (SQLite) ===') print("\n=== Batch Summary (SQLite) ===")
print(f'Database: {DB_PATH}') print(f"Database: {DB_PATH}")
ayam_tracked = {} ayam_tracked = {}
talenan_tracked = {} talenan_tracked = {}
if __name__ == '__main__': if __name__ == "__main__":
run() run()