diff --git a/counter_live_rknn.py b/counter_live_rknn.py index 4d98bd9..f4045f3 100644 --- a/counter_live_rknn.py +++ b/counter_live_rknn.py @@ -3,6 +3,7 @@ Edge production live counter — RTSP + YOLO RKNN + line crossing. Runs on RK3588 hardware with RKNN model (320×320 input). Replaces the Jetson/TensorRT variant. """ + import numpy as np import cv2 import csv @@ -13,68 +14,73 @@ from datetime import datetime from pathlib import Path from dotenv import load_dotenv + load_dotenv() from rknnlite.api import RKNNLite from batch_store import BatchStore # --- config (override via env / .env) --- -OUTPUT_DIR = os.getenv('OUTPUT_DIR', '/opt/jetson-counter') -DB_PATH = os.getenv('DB_PATH', f'{OUTPUT_DIR}/jetson_counter.db') -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') -MODEL_PATH = os.getenv('MODEL_PATH', '/opt/jetson-counter/yolo11n.rknn') -CAMERA_NAME = os.getenv('CAMERA_NAME', 'CC1') -OBJECT_LABEL = os.getenv('OBJECT_LABEL', 'ayam-potong') -CLASS_AYAM = os.getenv('CLASS_AYAM', 'ayam') -CLASS_TALENAN = os.getenv('CLASS_TALENAN', 'talenan') +OUTPUT_DIR = os.getenv("OUTPUT_DIR", "/opt/jetson-counter") +DB_PATH = os.getenv("DB_PATH", f"{OUTPUT_DIR}/jetson_counter.db") +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") +MODEL_PATH = os.getenv("MODEL_PATH", "/opt/jetson-counter/yolo11n.rknn") +CAMERA_NAME = os.getenv("CAMERA_NAME", "CC1") +OBJECT_LABEL = os.getenv("OBJECT_LABEL", "ayam-potong") +CLASS_AYAM = os.getenv("CLASS_AYAM", "ayam") +CLASS_TALENAN = os.getenv("CLASS_TALENAN", "talenan") -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')) -CROSS_DIRECTION = os.getenv('CROSS_DIRECTION', 'rtl').lower() +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")) +CROSS_DIRECTION = os.getenv("CROSS_DIRECTION", "rtl").lower() -IMGSZ = int(os.getenv('IMGSZ', '320')) -HALF = os.getenv('HALF', 'false').lower() == 'true' -CONF = float(os.getenv('CONF', '0.3')) -DEVICE = int(os.getenv('DEVICE', '0')) +IMGSZ = int(os.getenv("IMGSZ", "320")) +HALF = os.getenv("HALF", "false").lower() == "true" +CONF = float(os.getenv("CONF", "0.3")) +DEVICE = int(os.getenv("DEVICE", "0")) # RKNN-specific — core mask for NPU # 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 -NUM_CLASSES = int(os.getenv('NUM_CLASSES', '2')) -SCORE_SIGMOID = os.getenv('SCORE_SIGMOID', 'false').lower() == 'true' +NUM_CLASSES = int(os.getenv("NUM_CLASSES", "2")) +SCORE_SIGMOID = os.getenv("SCORE_SIGMOID", "false").lower() == "true" -DAILY_CUTOFF_TIME = os.getenv('DAILY_CUTOFF_TIME', '20:00') -BATCH_TIMEOUT_SECONDS = float(os.getenv('BATCH_TIMEOUT_SECONDS', '300')) -IGNORE_BATCH_LABEL_TIMEOUT = float(os.getenv('IGNORE_BATCH_LABEL_TIMEOUT_SECONDS', '30')) -MIN_OBJECT_PER_BATCH = int(os.getenv('MIN_OBJECT_PER_BATCH', '60')) -MIN_DURATION_PER_BATCH = int(os.getenv('MIN_DURATION_PER_BATCH', '60')) +DAILY_CUTOFF_TIME = os.getenv("DAILY_CUTOFF_TIME", "20:00") +BATCH_TIMEOUT_SECONDS = float(os.getenv("BATCH_TIMEOUT_SECONDS", "300")) +IGNORE_BATCH_LABEL_TIMEOUT = float( + os.getenv("IGNORE_BATCH_LABEL_TIMEOUT_SECONDS", "30") +) +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' -CROSS_CSV = os.getenv('CROSS_CSV', f'{OUTPUT_DIR}/batch_crossings.csv') +EXPORT_CSV = os.getenv("EXPORT_CSV", "true").lower() == "true" +CROSS_CSV = os.getenv("CROSS_CSV", f"{OUTPUT_DIR}/batch_crossings.csv") -WARMUP_FRAMES = int(os.getenv('WARMUP_FRAMES', '30')) -RECONNECT_DELAY_SEC = int(os.getenv('RECONNECT_DELAY_SEC', '3')) -MAX_RECONNECT_ATTEMPTS = int(os.getenv('MAX_RECONNECT_ATTEMPTS', '0')) -FLUSH_EVERY_N_FRAMES = int(os.getenv('FLUSH_EVERY_N_FRAMES', '100')) -TRACKED_PRUNE_SEC = int(os.getenv('TRACKED_PRUNE_SEC', '300')) -RECORD_VIDEO = os.getenv('RECORD_VIDEO', 'false').lower() == 'true' -VIDEO_SEGMENT_SEC = int(os.getenv('VIDEO_SEGMENT_SEC', '3600')) -OUTPUT_FPS = int(os.getenv('OUTPUT_FPS', '15')) +WARMUP_FRAMES = int(os.getenv("WARMUP_FRAMES", "30")) +RECONNECT_DELAY_SEC = int(os.getenv("RECONNECT_DELAY_SEC", "3")) +MAX_RECONNECT_ATTEMPTS = int(os.getenv("MAX_RECONNECT_ATTEMPTS", "0")) +FLUSH_EVERY_N_FRAMES = int(os.getenv("FLUSH_EVERY_N_FRAMES", "100")) +TRACKED_PRUNE_SEC = int(os.getenv("TRACKED_PRUNE_SEC", "300")) +RECORD_VIDEO = os.getenv("RECORD_VIDEO", "false").lower() == "true" +VIDEO_SEGMENT_SEC = int(os.getenv("VIDEO_SEGMENT_SEC", "3600")) +OUTPUT_FPS = int(os.getenv("OUTPUT_FPS", "15")) -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_QUALITY = int(os.getenv('LIVE_STREAM_QUALITY', '75')) -LIVE_STREAM_EVERY_N = int(os.getenv('LIVE_STREAM_EVERY_N', '2')) +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_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( - 'OPENCV_FFMPEG_CAPTURE_OPTIONS', - 'rtsp_transport;tcp|fflags;nobuffer|flags;low_delay', + "OPENCV_FFMPEG_CAPTURE_OPTIONS", + "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 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)] SK_COLORS = [ - (0, 255, 255), (0, 255, 255), (255, 0, 255), (255, 0, 255), - (0, 255, 0), (255, 255, 0), (0, 0, 255), (200, 200, 0), + (0, 255, 255), + (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) @@ -105,7 +117,7 @@ shutdown_requested = False def request_shutdown(signum, frame): global shutdown_requested shutdown_requested = True - print('\nShutdown requested — finishing current frame...') + print("\nShutdown requested — finishing current frame...") 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) # ============================================================================= + def _nms(boxes, scores, iou_thr=0.45): order = np.argsort(scores)[::-1] keep = [] @@ -132,7 +145,9 @@ def _nms(boxes, scores, iou_thr=0.45): h = np.maximum(0.0, yy2 - yy1) inter = w * h 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) order = order[1:][iou < iou_thr] return np.array(keep) @@ -164,12 +179,13 @@ def _compute_iou(box1, boxes2): # Simple IoU tracker (replaces bytetrack — same persist behaviour) # ============================================================================= + class SimpleTracker: def __init__(self, max_age=30, min_hits=1, iou_threshold=0.3): self.max_age = max_age self.min_hits = min_hits 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 def update(self, detections): @@ -177,33 +193,41 @@ class SimpleTracker: now = time.monotonic() for tid in self.tracks: - self.tracks[tid]['time_since_update'] += 1 + self.tracks[tid]["time_since_update"] += 1 matched_det = set() matched_track = set() - assignments = [] # (track_id, det_idx) + assignments = [] # (track_id, det_idx) det_to_track = {} # det_idx → track_id if detections and self.tracks: 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): _, 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)) - 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)) matched_track.add(track_ids[best_j]) matched_det.add(di) for tid, di in assignments: cx, box = detections[di] - self.tracks[tid]['cx'] = cx - self.tracks[tid]['box'] = box - self.tracks[tid]['hits'] += 1 - self.tracks[tid]['time_since_update'] = 0 - self.tracks[tid]['last_update'] = now + self.tracks[tid]["cx"] = cx + self.tracks[tid]["box"] = box + self.tracks[tid]["hits"] += 1 + self.tracks[tid]["time_since_update"] = 0 + self.tracks[tid]["last_update"] = now det_to_track[di] = tid for di, det in enumerate(detections): @@ -212,19 +236,27 @@ class SimpleTracker: new_id = self.next_id self.next_id += 1 self.tracks[new_id] = { - 'cx': cx, 'box': box, 'hits': 1, - 'time_since_update': 0, 'last_update': now, + "cx": cx, + "box": box, + "hits": 1, + "time_since_update": 0, + "last_update": now, } det_to_track[di] = new_id - stale = [tid for tid, t in self.tracks.items() - if t['time_since_update'] > self.max_age] + stale = [ + tid + for tid, t in self.tracks.items() + if t["time_since_update"] > self.max_age + ] for tid in stale: del self.tracks[tid] - track_map = {tid: self.tracks[tid]['cx'] - for tid in self.tracks - if self.tracks[tid]['hits'] >= self.min_hits} + track_map = { + tid: self.tracks[tid]["cx"] + for tid in self.tracks + if self.tracks[tid]["hits"] >= self.min_hits + } return track_map, det_to_track @@ -232,9 +264,19 @@ class SimpleTracker: # RKNN YOLO wrapper (detect output format: (1, 4+num_classes, N)) # ============================================================================= + class RKNNYOLO: - def __init__(self, model_path, core_mask=1, imgsz=320, conf=0.3, iou=0.45, - num_classes=2, num_keypoints=0, score_sigmoid=False): + def __init__( + 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.conf = conf self.iou = iou @@ -245,19 +287,20 @@ class RKNNYOLO: self.rknn = RKNNLite(verbose=False) ret = self.rknn.load_rknn(model_path) 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) 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: from rknnlite.api import RKNNLite as _RK + sdk_ver = self.rknn.get_sdk_version() - print(f'RKNN SDK version: {sdk_ver}') + print(f"RKNN SDK version: {sdk_ver}") except Exception: 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): """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) dy = (self.imgsz - nh) // 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) gains = np.array([scale, scale, dy, dx], dtype=np.float32) @@ -295,7 +338,7 @@ class RKNNYOLO: if out.shape[0] == self.num_classes + 4: 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() if self.score_sigmoid: @@ -303,12 +346,15 @@ class RKNNYOLO: else: cls_scores = cls_raw - 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[:, 1] + boxes_cxcywh[:, 3] / 2, - ], axis=1) + 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[:, 1] + boxes_cxcywh[:, 3] / 2, + ], + axis=1, + ) max_scores = cls_scores.max(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) for k in keep: j = idx[k] - detections.append({ - 'bbox': bboxes[j].tolist(), - 'score': float(scores[j]), - 'cls': int(clses[j]), - 'keypoints': None, - }) + detections.append( + { + "bbox": bboxes[j].tolist(), + "score": float(scores[j]), + "cls": int(clses[j]), + "keypoints": None, + } + ) return detections def release(self): @@ -365,41 +413,41 @@ def resolve_line_x(frame_width): def crossed_line(prev_cx, cx, line_x, direction=CROSS_DIRECTION): - if direction == 'ltr': + if direction == "ltr": 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 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): - if source.lower().startswith(('rtsp://', 'http://')): - os.environ['OPENCV_FFMPEG_CAPTURE_OPTIONS'] = RTSP_FFMPEG_OPTIONS + if source.lower().startswith(("rtsp://", "http://")): + os.environ["OPENCV_FFMPEG_CAPTURE_OPTIONS"] = RTSP_FFMPEG_OPTIONS cap = cv2.VideoCapture(source, cv2.CAP_FFMPEG) cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) return cap def warmup_stream(cap, n=WARMUP_FRAMES): - print('Warming up stream...') + print("Warming up stream...") for _ in range(n): cap.read() - print('Stream ready!') + print("Stream ready!") 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: def __init__(self, path, header): Path(path).parent.mkdir(parents=True, exist_ok=True) 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) if new_file: self.writer.writerow(header) @@ -424,8 +472,8 @@ class VideoSegmentWriter: self._open_next() def _segment_path(self): - ts = datetime.now().strftime('%Y%m%d_%H%M%S') - return str(self.output_dir / f'live_{ts}.mp4') + ts = datetime.now().strftime("%Y%m%d_%H%M%S") + return str(self.output_dir / f"live_{ts}.mp4") def _open_next(self): if self.writer is not None: @@ -433,7 +481,7 @@ class VideoSegmentWriter: path = self._segment_path() self.writer = open_video_writer(path, self.w, self.h, self.fps) self.segment_start = time.monotonic() - print(f'Recording segment: {path}') + print(f"Recording segment: {path}") def write(self, frame): if time.monotonic() - self.segment_start >= self.segment_sec: @@ -446,7 +494,9 @@ class VideoSegmentWriter: 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: 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) cv2.line(img, (line_x, y), (line_x, y_end), C_LINE_CORE, 2, cv2.LINE_AA) 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): @@ -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) pad = 14 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) - cv2.rectangle(img, (tx - pad, ty - th - pad), (tx + tw + pad, ty + pad // 2), C_LINE_CORE, 2) + overlay_rect( + 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) -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 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.putText(img, 'BATCH', (16, 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, 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) - 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) + cv2.putText( + img, "BATCH", (16, 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, + 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 + ) + 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): bar_h = 28 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): @@ -542,28 +687,39 @@ def draw_skeleton_bold(img, kpts): def draw_popups(img, popups, frame_idx): alive = [] for pop in popups: - age = frame_idx - pop['born'] + age = frame_idx - pop["born"] if age > POPUP_LIFETIME: continue alive.append(pop) 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)) - 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 def draw_batch_banner(img, w, batch_num, pulse_remaining): if pulse_remaining <= 0: return - text = f'NEW BATCH {batch_num}' + text = f"NEW BATCH {batch_num}" font = cv2.FONT_HERSHEY_SIMPLEX (tw, th), _ = cv2.getTextSize(text, font, 0.8, 2) x1, y1 = w // 2 - tw // 2 - 16, 62 x2, y2 = w // 2 + tw // 2 + 16, 62 + th + 20 overlay_rect(img, x1, y1, x2, y2, C_PANEL, alpha=0.7) 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): @@ -573,11 +729,13 @@ def connect_stream(source, warmup=WARMUP_FRAMES): if not cap.isOpened(): attempts += 1 if MAX_RECONNECT_ATTEMPTS and attempts >= MAX_RECONNECT_ATTEMPTS: - raise RuntimeError(f'Cannot open source after {attempts} attempts: {source}') - print(f'Cannot open source, retry in {RECONNECT_DELAY_SEC}s...') + raise RuntimeError( + f"Cannot open source after {attempts} attempts: {source}" + ) + print(f"Cannot open source, retry in {RECONNECT_DELAY_SEC}s...") time.sleep(RECONNECT_DELAY_SEC) continue - if warmup > 0 and source.lower().startswith(('rtsp://', 'http://')): + if warmup > 0 and source.lower().startswith(("rtsp://", "http://")): warmup_stream(cap, warmup) w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) @@ -592,6 +750,7 @@ def connect_stream(source, warmup=WARMUP_FRAMES): # Main loop # ============================================================================= + def run(): global shutdown_requested @@ -605,13 +764,15 @@ def run(): ignore_batch_label_timeout=IGNORE_BATCH_LABEL_TIMEOUT, min_object_per_batch=MIN_OBJECT_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() cross_logger = None 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 model = RKNNYOLO( @@ -627,8 +788,8 @@ def run(): # 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 CLASS_IDS = { - os.getenv('CLASS_AYAM', 'ayam'): 0, - os.getenv('CLASS_TALENAN', 'talenan'): 1, + os.getenv("CLASS_AYAM", "ayam"): 0, + os.getenv("CLASS_TALENAN", "talenan"): 1, } ayam_cls = CLASS_IDS[CLASS_AYAM] talenan_cls = CLASS_IDS[CLASS_TALENAN] @@ -655,10 +816,12 @@ def run(): return line_x = resolve_line_x(w) - print(f'RKNN counter | {w}x{h} @ {fps}fps | line x={line_x} | cross={CROSS_DIRECTION}') - print(f'Model: {MODEL_PATH} | imgsz={IMGSZ} | core_mask={CORE_MASK}') - print(f'DB: {DB_PATH}') - print(f'State: {STATE_FILE}') + print( + f"RKNN counter | {w}x{h} @ {fps}fps | line x={line_x} | cross={CROSS_DIRECTION}" + ) + print(f"Model: {MODEL_PATH} | imgsz={IMGSZ} | core_mask={CORE_MASK}") + print(f"DB: {DB_PATH}") + print(f"State: {STATE_FILE}") if RECORD_VIDEO: video_writer = VideoSegmentWriter(OUTPUT_DIR, w, h, fps, VIDEO_SEGMENT_SEC) @@ -671,7 +834,9 @@ def run(): if not IS_LIVE: break 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() time.sleep(RECONNECT_DELAY_SEC) cap, w, h, fps = connect_stream(SOURCE) @@ -691,12 +856,12 @@ def run(): if detections: ayam_dets = [] # list of (cx, xywh_box) talenan_dets = [] - ayam_kpts_map = {} # det_idx → keypoints + ayam_kpts_map = {} # det_idx → keypoints talenan_kpts_map = {} for di, det in enumerate(detections): - bbox = det['bbox'] - cls_id = det['cls'] + bbox = det["bbox"] + cls_id = det["cls"] cx = (bbox[0] + bbox[2]) / 2.0 x1, y1, x2, y2 = bbox wb, hb = x2 - x1, y2 - y1 @@ -704,12 +869,12 @@ def run(): if cls_id == talenan_cls: talenan_dets.append((cx, box_cxcywh)) - if det['keypoints'] is not None: - talenan_kpts_map[len(talenan_dets) - 1] = det['keypoints'] + if det["keypoints"] is not None: + talenan_kpts_map[len(talenan_dets) - 1] = det["keypoints"] elif cls_id == ayam_cls: ayam_dets.append((cx, box_cxcywh)) - if det['keypoints'] is not None: - ayam_kpts_map[len(ayam_dets) - 1] = det['keypoints'] + if det["keypoints"] is not None: + ayam_kpts_map[len(ayam_dets) - 1] = det["keypoints"] # Track ayam — returns (track_id → cx, detection_idx → track_id) 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_tracked: 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) if store.record_talenan_crossing(tid): batch_closed_frame = True talenan_cross_flash[tid] = CROSS_FLASH_FRAMES - popups.append({ - 'x': int(cx) - 20, - 'y': int(box[1]), - 'born': frame_idx, - 'text': 'BATCH CLOSED', - }) + popups.append( + { + "x": int(cx) - 20, + "y": int(box[1]), + "born": frame_idx, + "text": "BATCH CLOSED", + } + ) talenan_tracked[tid] = (cx, mono) # Process ayam crossings @@ -746,24 +916,33 @@ def run(): if tid in ayam_tracker.tracks: if tid in ayam_tracked: 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) _, started_new = store.record_ayam_crossing(tid) if cross_logger: - cross_logger.write_row([ - store.current_batch_number, frame_idx, - datetime.now().isoformat(), tid, - ]) + cross_logger.write_row( + [ + store.current_batch_number, + frame_idx, + datetime.now().isoformat(), + tid, + ] + ) ayam_crossed_frame = True if started_new: batch_started_frame = True ayam_cross_flash[tid] = CROSS_FLASH_FRAMES - popups.append({ - 'x': int(cx) - 12, - 'y': int(box[1]), - 'born': frame_idx, - 'text': '+1', - }) + popups.append( + { + "x": int(cx) - 12, + "y": int(box[1]), + "born": frame_idx, + "text": "+1", + } + ) ayam_tracked[tid] = (cx, mono) # Draw talenan @@ -778,7 +957,7 @@ def run(): flash = talenan_cross_flash.get(tid, 0) color = C_GREEN if flash > 0 else C_TALENAN_BOX 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 for di, (cx, box) in enumerate(ayam_dets): @@ -792,7 +971,7 @@ def run(): flash = ayam_cross_flash.get(tid, 0) color = C_GREEN if flash > 0 else C_AYAM_BOX 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) if kpts is not None: draw_skeleton_bold(frame, kpts) @@ -812,9 +991,19 @@ def run(): draw_elegant_counting_line(frame, line_x, h, line_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_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) 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: try: Path(LIVE_STREAM_FRAME_PATH).parent.mkdir(parents=True, exist_ok=True) - _, jpeg = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, LIVE_STREAM_QUALITY]) - with open(LIVE_STREAM_FRAME_PATH, 'wb') as f: + _, jpeg = cv2.imencode( + ".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, LIVE_STREAM_QUALITY] + ) + with open(LIVE_STREAM_FRAME_PATH, "wb") as f: f.write(jpeg.tobytes()) except Exception: pass @@ -841,8 +1032,8 @@ def run(): frame_idx += 1 if frame_idx % FLUSH_EVERY_N_FRAMES == 0: print( - f'[{now_str()}] Frame {frame_idx} | Batch {batch_num}: {batch_count} ' - f'| Total: {display_total} | Uptime {elapsed / 3600:.2f}h' + f"[{now_str()}] Frame {frame_idx} | Batch {batch_num}: {batch_count} " + f"| Total: {display_total} | Uptime {elapsed / 3600:.2f}h" ) prune_stale_tracks(ayam_tracked, mono) prune_stale_tracks(talenan_tracked, mono) @@ -855,8 +1046,8 @@ def run(): model.release() store.shutdown() - print('\n=== Batch Summary (SQLite) ===') - print(f'Database: {DB_PATH}') + print("\n=== Batch Summary (SQLite) ===") + print(f"Database: {DB_PATH}") # Tracked state dicts: track_id → (cx, monotonic_time) @@ -864,5 +1055,5 @@ ayam_tracked = {} talenan_tracked = {} -if __name__ == '__main__': +if __name__ == "__main__": run() diff --git a/counter_live_rknn_bytetrack.py b/counter_live_rknn_bytetrack.py index f905e17..1ad5e62 100644 --- a/counter_live_rknn_bytetrack.py +++ b/counter_live_rknn_bytetrack.py @@ -3,6 +3,7 @@ Edge production live counter — RTSP + YOLO RKNN + ByteTrack + line crossing. Runs on RK3588 hardware with RKNN model (320×320 input). Uses ByteTrack (Kalman filter + two-stage IoU association) for tracking. """ + import numpy as np import cv2 import csv @@ -13,74 +14,79 @@ from datetime import datetime from pathlib import Path from dotenv import load_dotenv + load_dotenv() from rknnlite.api import RKNNLite from batch_store import BatchStore # --- config (override via env / .env) --- -OUTPUT_DIR = os.getenv('OUTPUT_DIR', '/opt/jetson-counter') -DB_PATH = os.getenv('DB_PATH', f'{OUTPUT_DIR}/jetson_counter.db') -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') -MODEL_PATH = os.getenv('MODEL_PATH', '/opt/jetson-counter/yolo11n.rknn') -CAMERA_NAME = os.getenv('CAMERA_NAME', 'CC1') -OBJECT_LABEL = os.getenv('OBJECT_LABEL', 'ayam-potong') -CLASS_AYAM = os.getenv('CLASS_AYAM', 'ayam') -CLASS_TALENAN = os.getenv('CLASS_TALENAN', 'talenan') +OUTPUT_DIR = os.getenv("OUTPUT_DIR", "/opt/jetson-counter") +DB_PATH = os.getenv("DB_PATH", f"{OUTPUT_DIR}/jetson_counter.db") +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") +MODEL_PATH = os.getenv("MODEL_PATH", "/opt/jetson-counter/yolo11n.rknn") +CAMERA_NAME = os.getenv("CAMERA_NAME", "CC1") +OBJECT_LABEL = os.getenv("OBJECT_LABEL", "ayam-potong") +CLASS_AYAM = os.getenv("CLASS_AYAM", "ayam") +CLASS_TALENAN = os.getenv("CLASS_TALENAN", "talenan") -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')) -CROSS_DIRECTION = os.getenv('CROSS_DIRECTION', 'rtl').lower() +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")) +CROSS_DIRECTION = os.getenv("CROSS_DIRECTION", "rtl").lower() -IMGSZ = int(os.getenv('IMGSZ', '320')) -HALF = os.getenv('HALF', 'false').lower() == 'true' -CONF = float(os.getenv('CONF', '0.3')) -DEVICE = int(os.getenv('DEVICE', '0')) +IMGSZ = int(os.getenv("IMGSZ", "320")) +HALF = os.getenv("HALF", "false").lower() == "true" +CONF = float(os.getenv("CONF", "0.3")) +DEVICE = int(os.getenv("DEVICE", "0")) # RKNN NPU core mask -CORE_MASK = int(os.getenv('CORE_MASK', '1')) +CORE_MASK = int(os.getenv("CORE_MASK", "1")) # YOLO decoder config -NUM_CLASSES = int(os.getenv('NUM_CLASSES', '2')) -SCORE_SIGMOID = os.getenv('SCORE_SIGMOID', 'false').lower() == 'true' +NUM_CLASSES = int(os.getenv("NUM_CLASSES", "2")) +SCORE_SIGMOID = os.getenv("SCORE_SIGMOID", "false").lower() == "true" # ByteTrack settings -TRACK_HIGH_THRESH = float(os.getenv('TRACK_HIGH_THRESH', '0.5')) -TRACK_LOW_THRESH = float(os.getenv('TRACK_LOW_THRESH', '0.1')) -TRACK_MATCH_THRESH = float(os.getenv('TRACK_MATCH_THRESH', '0.8')) -TRACK_BUFFER = int(os.getenv('TRACK_BUFFER', '30')) -TRACK_MIN_HITS = int(os.getenv('TRACK_MIN_HITS', '3')) +TRACK_HIGH_THRESH = float(os.getenv("TRACK_HIGH_THRESH", "0.5")) +TRACK_LOW_THRESH = float(os.getenv("TRACK_LOW_THRESH", "0.1")) +TRACK_MATCH_THRESH = float(os.getenv("TRACK_MATCH_THRESH", "0.8")) +TRACK_BUFFER = int(os.getenv("TRACK_BUFFER", "30")) +TRACK_MIN_HITS = int(os.getenv("TRACK_MIN_HITS", "3")) -DAILY_CUTOFF_TIME = os.getenv('DAILY_CUTOFF_TIME', '20:00') -BATCH_TIMEOUT_SECONDS = float(os.getenv('BATCH_TIMEOUT_SECONDS', '300')) -IGNORE_BATCH_LABEL_TIMEOUT = float(os.getenv('IGNORE_BATCH_LABEL_TIMEOUT_SECONDS', '30')) -MIN_OBJECT_PER_BATCH = int(os.getenv('MIN_OBJECT_PER_BATCH', '60')) -MIN_DURATION_PER_BATCH = int(os.getenv('MIN_DURATION_PER_BATCH', '60')) +DAILY_CUTOFF_TIME = os.getenv("DAILY_CUTOFF_TIME", "20:00") +BATCH_TIMEOUT_SECONDS = float(os.getenv("BATCH_TIMEOUT_SECONDS", "300")) +IGNORE_BATCH_LABEL_TIMEOUT = float( + os.getenv("IGNORE_BATCH_LABEL_TIMEOUT_SECONDS", "30") +) +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' -CROSS_CSV = os.getenv('CROSS_CSV', f'{OUTPUT_DIR}/batch_crossings.csv') +EXPORT_CSV = os.getenv("EXPORT_CSV", "true").lower() == "true" +CROSS_CSV = os.getenv("CROSS_CSV", f"{OUTPUT_DIR}/batch_crossings.csv") -WARMUP_FRAMES = int(os.getenv('WARMUP_FRAMES', '30')) -RECONNECT_DELAY_SEC = int(os.getenv('RECONNECT_DELAY_SEC', '3')) -MAX_RECONNECT_ATTEMPTS = int(os.getenv('MAX_RECONNECT_ATTEMPTS', '0')) -FLUSH_EVERY_N_FRAMES = int(os.getenv('FLUSH_EVERY_N_FRAMES', '100')) -TRACKED_PRUNE_SEC = int(os.getenv('TRACKED_PRUNE_SEC', '300')) -RECORD_VIDEO = os.getenv('RECORD_VIDEO', 'false').lower() == 'true' -VIDEO_SEGMENT_SEC = int(os.getenv('VIDEO_SEGMENT_SEC', '3600')) -OUTPUT_FPS = int(os.getenv('OUTPUT_FPS', '15')) +WARMUP_FRAMES = int(os.getenv("WARMUP_FRAMES", "30")) +RECONNECT_DELAY_SEC = int(os.getenv("RECONNECT_DELAY_SEC", "3")) +MAX_RECONNECT_ATTEMPTS = int(os.getenv("MAX_RECONNECT_ATTEMPTS", "0")) +FLUSH_EVERY_N_FRAMES = int(os.getenv("FLUSH_EVERY_N_FRAMES", "100")) +TRACKED_PRUNE_SEC = int(os.getenv("TRACKED_PRUNE_SEC", "300")) +RECORD_VIDEO = os.getenv("RECORD_VIDEO", "false").lower() == "true" +VIDEO_SEGMENT_SEC = int(os.getenv("VIDEO_SEGMENT_SEC", "3600")) +OUTPUT_FPS = int(os.getenv("OUTPUT_FPS", "15")) -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_QUALITY = int(os.getenv('LIVE_STREAM_QUALITY', '75')) -LIVE_STREAM_EVERY_N = int(os.getenv('LIVE_STREAM_EVERY_N', '2')) +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_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( - 'OPENCV_FFMPEG_CAPTURE_OPTIONS', - 'rtsp_transport;tcp|fflags;nobuffer|flags;low_delay', + "OPENCV_FFMPEG_CAPTURE_OPTIONS", + "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 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)] SK_COLORS = [ - (0, 255, 255), (0, 255, 255), (255, 0, 255), (255, 0, 255), - (0, 255, 0), (255, 255, 0), (0, 0, 255), (200, 200, 0), + (0, 255, 255), + (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) @@ -111,7 +123,7 @@ shutdown_requested = False def request_shutdown(signum, frame): global shutdown_requested shutdown_requested = True - print('\nShutdown requested — finishing current frame...') + print("\nShutdown requested — finishing current frame...") 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) # ============================================================================= + def _nms(boxes, scores, iou_thr=0.45): order = np.argsort(scores)[::-1] keep = [] @@ -138,7 +151,9 @@ def _nms(boxes, scores, iou_thr=0.45): h = np.maximum(0.0, yy2 - yy1) inter = w * h 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) order = order[1:][iou < iou_thr] return np.array(keep) @@ -148,6 +163,7 @@ def _nms(boxes, scores, iou_thr=0.45): # IoU helpers (xyxy format) # ============================================================================= + def _ious_xyxy(boxes_a, boxes_b): """Pairwise IoU: (N,4) vs (M,4) → (N,M) matrix.""" 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) # ============================================================================= + class KalmanBoxTracker: count = 0 @@ -242,6 +259,7 @@ class KalmanBoxTracker: class _KalmanFilter: """8-state constant-velocity Kalman filter for bounding box tracking.""" + def __init__(self): ndim, dt = 4, 1.0 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 def update(self, z): - R = np.diag(np.square([ - self._std_weight_position * z[2], - self._std_weight_position * z[3], - self._std_weight_position * z[2], - self._std_weight_position * z[3], - ])) + R = np.diag( + np.square( + [ + self._std_weight_position * z[2], + self._std_weight_position * z[3], + self._std_weight_position * z[2], + self._std_weight_position * z[3], + ] + ) + ) H = self.update_mat S = H @ self.P @ H.T + R @@ -295,11 +317,18 @@ class _KalmanFilter: # ByteTrack multi-object tracker # ============================================================================= + class ByteTracker: """ByteTrack: two-stage association with Kalman filter prediction.""" - def __init__(self, track_high_thresh=0.5, track_low_thresh=0.1, - match_thresh=0.8, track_buffer=30, min_hits=3): + def __init__( + 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.low_thresh = track_low_thresh self.match_thresh = match_thresh @@ -367,7 +396,9 @@ class ByteTracker: tracked_map[track_pool[ti].track_id] = track_pool[ti].get_cx() 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 --- low_idx = np.where(is_low)[0] @@ -383,10 +414,14 @@ class ByteTracker: det_global = int(low_idx[di]) pool_idx = unmatched_tracks[uti] 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) 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 --- for ti, trk in enumerate(track_pool): @@ -432,9 +467,19 @@ class ByteTracker: # RKNN YOLO wrapper (detect output format: (1, 4+num_classes, N)) # ============================================================================= + class RKNNYOLO: - def __init__(self, model_path, core_mask=1, imgsz=320, conf=0.3, iou=0.45, - num_classes=2, num_keypoints=0, score_sigmoid=False): + def __init__( + 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.conf = conf self.iou = iou @@ -445,18 +490,18 @@ class RKNNYOLO: self.rknn = RKNNLite(verbose=False) ret = self.rknn.load_rknn(model_path) 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) 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: sdk_ver = self.rknn.get_sdk_version() - print(f'RKNN SDK version: {sdk_ver}') + print(f"RKNN SDK version: {sdk_ver}") except Exception: 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): h0, w0 = frame.shape[:2] @@ -467,7 +512,7 @@ class RKNNYOLO: letterbox = np.full((self.imgsz, self.imgsz, 3), 114, dtype=np.uint8) dy = (self.imgsz - nh) // 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) gains = np.array([scale, scale, dy, dx], dtype=np.float32) @@ -500,12 +545,15 @@ class RKNNYOLO: else: cls_scores = cls_raw - 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[:, 1] + boxes_cxcywh[:, 3] / 2, - ], axis=1) + 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[:, 1] + boxes_cxcywh[:, 3] / 2, + ], + axis=1, + ) max_scores = cls_scores.max(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) for k in keep: j = idx[k] - detections.append({ - 'bbox': bboxes[j].tolist(), - 'score': float(scores[j]), - 'cls': int(clses[j]), - 'keypoints': None, - }) + detections.append( + { + "bbox": bboxes[j].tolist(), + "score": float(scores[j]), + "cls": int(clses[j]), + "keypoints": None, + } + ) return detections def release(self): @@ -552,6 +602,7 @@ class RKNNYOLO: # Drawing helpers # ============================================================================= + def resolve_line_x(frame_width): if LINE_X is not None: return LINE_X @@ -561,41 +612,41 @@ def resolve_line_x(frame_width): def crossed_line(prev_cx, cx, line_x, direction=CROSS_DIRECTION): - if direction == 'ltr': + if direction == "ltr": 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 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): - if source.lower().startswith(('rtsp://', 'http://')): - os.environ['OPENCV_FFMPEG_CAPTURE_OPTIONS'] = RTSP_FFMPEG_OPTIONS + if source.lower().startswith(("rtsp://", "http://")): + os.environ["OPENCV_FFMPEG_CAPTURE_OPTIONS"] = RTSP_FFMPEG_OPTIONS cap = cv2.VideoCapture(source, cv2.CAP_FFMPEG) cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) return cap def warmup_stream(cap, n=WARMUP_FRAMES): - print('Warming up stream...') + print("Warming up stream...") for _ in range(n): cap.read() - print('Stream ready!') + print("Stream ready!") 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: def __init__(self, path, header): Path(path).parent.mkdir(parents=True, exist_ok=True) 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) if new_file: self.writer.writerow(header) @@ -620,8 +671,8 @@ class VideoSegmentWriter: self._open_next() def _segment_path(self): - ts = datetime.now().strftime('%Y%m%d_%H%M%S') - return str(self.output_dir / f'live_{ts}.mp4') + ts = datetime.now().strftime("%Y%m%d_%H%M%S") + return str(self.output_dir / f"live_{ts}.mp4") def _open_next(self): if self.writer is not None: @@ -629,7 +680,7 @@ class VideoSegmentWriter: path = self._segment_path() self.writer = open_video_writer(path, self.w, self.h, self.fps) self.segment_start = time.monotonic() - print(f'Recording segment: {path}') + print(f"Recording segment: {path}") def write(self, frame): if time.monotonic() - self.segment_start >= self.segment_sec: @@ -642,7 +693,9 @@ class VideoSegmentWriter: 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: 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) cv2.line(img, (line_x, y), (line_x, y_end), C_LINE_CORE, 2, cv2.LINE_AA) 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): @@ -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) pad = 14 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) - cv2.rectangle(img, (tx - pad, ty - th - pad), (tx + tw + pad, ty + pad // 2), C_LINE_CORE, 2) + overlay_rect( + 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) -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 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.putText(img, 'BATCH', (16, 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, 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) - 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) + cv2.putText( + img, "BATCH", (16, 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, + 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 + ) + 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): bar_h = 28 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): @@ -738,28 +886,39 @@ def draw_skeleton_bold(img, kpts): def draw_popups(img, popups, frame_idx): alive = [] for pop in popups: - age = frame_idx - pop['born'] + age = frame_idx - pop["born"] if age > POPUP_LIFETIME: continue alive.append(pop) 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)) - 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 def draw_batch_banner(img, w, batch_num, pulse_remaining): if pulse_remaining <= 0: return - text = f'NEW BATCH {batch_num}' + text = f"NEW BATCH {batch_num}" font = cv2.FONT_HERSHEY_SIMPLEX (tw, th), _ = cv2.getTextSize(text, font, 0.8, 2) x1, y1 = w // 2 - tw // 2 - 16, 62 x2, y2 = w // 2 + tw // 2 + 16, 62 + th + 20 overlay_rect(img, x1, y1, x2, y2, C_PANEL, alpha=0.7) 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): @@ -769,11 +928,13 @@ def connect_stream(source, warmup=WARMUP_FRAMES): if not cap.isOpened(): attempts += 1 if MAX_RECONNECT_ATTEMPTS and attempts >= MAX_RECONNECT_ATTEMPTS: - raise RuntimeError(f'Cannot open source after {attempts} attempts: {source}') - print(f'Cannot open source, retry in {RECONNECT_DELAY_SEC}s...') + raise RuntimeError( + f"Cannot open source after {attempts} attempts: {source}" + ) + print(f"Cannot open source, retry in {RECONNECT_DELAY_SEC}s...") time.sleep(RECONNECT_DELAY_SEC) continue - if warmup > 0 and source.lower().startswith(('rtsp://', 'http://')): + if warmup > 0 and source.lower().startswith(("rtsp://", "http://")): warmup_stream(cap, warmup) w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) @@ -788,6 +949,7 @@ def connect_stream(source, warmup=WARMUP_FRAMES): # Main loop # ============================================================================= + def run(): global shutdown_requested @@ -801,13 +963,15 @@ def run(): ignore_batch_label_timeout=IGNORE_BATCH_LABEL_TIMEOUT, min_object_per_batch=MIN_OBJECT_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() cross_logger = None 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_path=MODEL_PATH, @@ -819,8 +983,8 @@ def run(): ) CLASS_IDS = { - os.getenv('CLASS_AYAM', 'ayam'): 0, - os.getenv('CLASS_TALENAN', 'talenan'): 1, + os.getenv("CLASS_AYAM", "ayam"): 0, + os.getenv("CLASS_TALENAN", "talenan"): 1, } ayam_cls = CLASS_IDS[CLASS_AYAM] talenan_cls = CLASS_IDS[CLASS_TALENAN] @@ -859,12 +1023,16 @@ def run(): return line_x = resolve_line_x(w) - print(f'RKNN+ByteTrack counter | {w}x{h} @ {fps}fps | line x={line_x} | cross={CROSS_DIRECTION}') - print(f'Model: {MODEL_PATH} | imgsz={IMGSZ} | core_mask={CORE_MASK}') - print(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}') + print( + f"RKNN+ByteTrack counter | {w}x{h} @ {fps}fps | line x={line_x} | cross={CROSS_DIRECTION}" + ) + print(f"Model: {MODEL_PATH} | imgsz={IMGSZ} | core_mask={CORE_MASK}") + print( + 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: video_writer = VideoSegmentWriter(OUTPUT_DIR, w, h, fps, VIDEO_SEGMENT_SEC) @@ -877,7 +1045,9 @@ def run(): if not IS_LIVE: break 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() time.sleep(RECONNECT_DELAY_SEC) cap, w, h, fps = connect_stream(SOURCE) @@ -905,10 +1075,10 @@ def run(): talenan_cx_list = [] for det in detections: - bbox = det['bbox'] - score = det['score'] - cls_id = det['cls'] - kpts = det['keypoints'] + bbox = det["bbox"] + score = det["score"] + cls_id = det["cls"] + kpts = det["keypoints"] cx = (bbox[0] + bbox[2]) / 2.0 if cls_id == talenan_cls: @@ -928,9 +1098,11 @@ def run(): talenan_scores = np.array(talenan_scores, dtype=np.float32) ayam_track_map, ayam_det_to_track, ayam_lost_map = ayam_tracker.update( - ayam_boxes_xyxy, ayam_scores) - talenan_track_map, talenan_det_to_track, talenan_lost_map = talenan_tracker.update( - talenan_boxes_xyxy, talenan_scores) + ayam_boxes_xyxy, ayam_scores + ) + talenan_track_map, talenan_det_to_track, talenan_lost_map = ( + talenan_tracker.update(talenan_boxes_xyxy, talenan_scores) + ) # Process talenan crossings for di in range(len(talenan_boxes_xyxy)): @@ -942,17 +1114,22 @@ def run(): if tid in talenan_tracked: 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) if store.record_talenan_crossing(tid): batch_closed_frame = True talenan_cross_flash[tid] = CROSS_FLASH_FRAMES - popups.append({ - 'x': int(cx) - 20, - 'y': int((bbox[1] + bbox[3]) / 2), - 'born': frame_idx, - 'text': 'BATCH CLOSED', - }) + popups.append( + { + "x": int(cx) - 20, + "y": int((bbox[1] + bbox[3]) / 2), + "born": frame_idx, + "text": "BATCH CLOSED", + } + ) talenan_tracked[tid] = (cx, mono) # Process ayam crossings (including lost tracks for line-cross continuity) @@ -964,24 +1141,36 @@ def run(): if tid in ayam_tracked: 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) _, started_new = store.record_ayam_crossing(tid) if cross_logger: - cross_logger.write_row([ - store.current_batch_number, frame_idx, - datetime.now().isoformat(), tid, - ]) + cross_logger.write_row( + [ + store.current_batch_number, + frame_idx, + datetime.now().isoformat(), + tid, + ] + ) ayam_crossed_frame = True if started_new: batch_started_frame = True ayam_cross_flash[tid] = CROSS_FLASH_FRAMES - popups.append({ - 'x': int(cx) - 12, - 'y': int((ayam_boxes_xyxy[di][1] + ayam_boxes_xyxy[di][3]) / 2), - 'born': frame_idx, - 'text': '+1', - }) + popups.append( + { + "x": int(cx) - 12, + "y": int( + (ayam_boxes_xyxy[di][1] + ayam_boxes_xyxy[di][3]) + / 2 + ), + "born": frame_idx, + "text": "+1", + } + ) ayam_tracked[tid] = (cx, mono) # Also track lost tracks for line-crossing continuity @@ -999,7 +1188,7 @@ def run(): flash = talenan_cross_flash.get(tid, 0) color = C_GREEN if flash > 0 else C_TALENAN_BOX 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 for di in range(len(ayam_boxes_xyxy)): @@ -1011,7 +1200,7 @@ def run(): flash = ayam_cross_flash.get(tid, 0) color = C_GREEN if flash > 0 else C_AYAM_BOX 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 if kpts is not None: draw_skeleton_bold(frame, kpts) @@ -1031,9 +1220,21 @@ def run(): draw_elegant_counting_line(frame, line_x, h, line_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_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) 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: try: Path(LIVE_STREAM_FRAME_PATH).parent.mkdir(parents=True, exist_ok=True) - _, jpeg = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, LIVE_STREAM_QUALITY]) - with open(LIVE_STREAM_FRAME_PATH, 'wb') as f: + _, jpeg = cv2.imencode( + ".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, LIVE_STREAM_QUALITY] + ) + with open(LIVE_STREAM_FRAME_PATH, "wb") as f: f.write(jpeg.tobytes()) except Exception: pass @@ -1060,8 +1263,8 @@ def run(): frame_idx += 1 if frame_idx % FLUSH_EVERY_N_FRAMES == 0: print( - f'[{now_str()}] Frame {frame_idx} | Batch {batch_num}: {batch_count} ' - f'| Total: {display_total} | Uptime {elapsed / 3600:.2f}h' + f"[{now_str()}] Frame {frame_idx} | Batch {batch_num}: {batch_count} " + f"| Total: {display_total} | Uptime {elapsed / 3600:.2f}h" ) prune_stale_tracks(ayam_tracked, mono) prune_stale_tracks(talenan_tracked, mono) @@ -1074,13 +1277,13 @@ def run(): model.release() store.shutdown() - print('\n=== Batch Summary (SQLite) ===') - print(f'Database: {DB_PATH}') + print("\n=== Batch Summary (SQLite) ===") + print(f"Database: {DB_PATH}") ayam_tracked = {} talenan_tracked = {} -if __name__ == '__main__': +if __name__ == "__main__": run()