""" 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 import os import signal import time from collections import deque 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") 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")) # RKNN NPU core mask 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" # ByteTrack settings # Ayam index 0 TRACK_HIGH_THRESH_0 = float(os.getenv("TRACK_HIGH_THRESH_0", "0.5")) TRACK_LOW_THRESH_0 = float(os.getenv("TRACK_LOW_THRESH_0", "0.1")) TRACK_MATCH_THRESH_0 = float(os.getenv("TRACK_MATCH_THRESH_0", "0.8")) TRACK_BUFFER_0 = int(os.getenv("TRACK_BUFFER_0", "30")) TRACK_MIN_HITS_0 = int(os.getenv("TRACK_MIN_HITS_0", "3")) # Talenan index 1 TRACK_HIGH_THRESH_1 = float(os.getenv("TRACK_HIGH_THRESH_1", "0.5")) TRACK_LOW_THRESH_1 = float(os.getenv("TRACK_LOW_THRESH_1", "0.1")) TRACK_MATCH_THRESH_1 = float(os.getenv("TRACK_MATCH_THRESH_1", "0.6")) TRACK_BUFFER_1 = int(os.getenv("TRACK_BUFFER_1", "30")) TRACK_MIN_HITS_1 = int(os.getenv("TRACK_MIN_HITS_1", "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")) EXPORT_CSV = os.getenv("EXPORT_CSV", "true").lower() == "true" CROSS_CSV = os.getenv("CROSS_CSV", f"{OUTPUT_DIR}/batch_crossings.csv") RATE_WINDOW_SEC = int(os.getenv("RATE_WINDOW_SEC", "60")) 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")) RTSP_FFMPEG_OPTIONS = os.getenv( "OPENCV_FFMPEG_CAPTURE_OPTIONS", "rtsp_transport;tcp|fflags;nobuffer|flags;low_delay", ) IS_LIVE = SOURCE.lower().startswith(("rtsp://", "http://")) MOTION_DETECTION_ENABLED = os.getenv("MOTION_DETECTION_ENABLED", "false").lower() == "true" MOTION_THRESHOLD = float(os.getenv("MOTION_THRESHOLD", "5.0")) CROSS_FLASH_FRAMES = 12 POPUP_LIFETIME = 20 LINE_PULSE_FRAMES = 12 COUNT_PULSE_FRAMES = 15 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), ] C_PANEL = (28, 24, 18) C_BORDER = (90, 85, 75) C_ACCENT = (255, 200, 60) C_GREEN = (80, 220, 100) C_TEXT = (235, 235, 235) C_MUTED = (150, 150, 150) C_AYAM_BOX = (0, 165, 255) C_TALENAN_BOX = (220, 120, 60) C_LINE_CORE = (180, 220, 255) C_LINE_GLOW = (100, 160, 220) shutdown_requested = False def request_shutdown(signum, frame): global shutdown_requested shutdown_requested = True print("\nShutdown requested — finishing current frame...") signal.signal(signal.SIGINT, request_shutdown) 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 = [] while len(order) > 0: idx = order[0] keep.append(idx) if len(order) == 1: break xx1 = np.maximum(boxes[idx, 0], boxes[order[1:], 0]) yy1 = np.maximum(boxes[idx, 1], boxes[order[1:], 1]) xx2 = np.minimum(boxes[idx, 2], boxes[order[1:], 2]) yy2 = np.minimum(boxes[idx, 3], boxes[order[1:], 3]) w = np.maximum(0.0, xx2 - xx1) 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] ) iou = inter / (area_i + area_o - inter + 1e-16) order = order[1:][iou < iou_thr] return np.array(keep) # ============================================================================= # 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) if n == 0 or m == 0: return np.zeros((n, m), dtype=np.float32) xx1 = np.maximum(boxes_a[:, None, 0], boxes_b[None, :, 0]) yy1 = np.maximum(boxes_a[:, None, 1], boxes_b[None, :, 1]) xx2 = np.minimum(boxes_a[:, None, 2], boxes_b[None, :, 2]) yy2 = np.minimum(boxes_a[:, None, 3], boxes_b[None, :, 3]) iw = np.maximum(0.0, xx2 - xx1) ih = np.maximum(0.0, yy2 - yy1) inter = iw * ih area_a = (boxes_a[:, 2] - boxes_a[:, 0]) * (boxes_a[:, 3] - boxes_a[:, 1]) area_b = (boxes_b[:, 2] - boxes_b[:, 0]) * (boxes_b[:, 3] - boxes_b[:, 1]) return inter / (area_a[:, None] + area_b[None, :] - inter + 1e-16) def _greedy_match(cost_matrix, threshold=0.3): """Greedy linear assignment. Returns pairs (row_idx, col_idx).""" if cost_matrix.size == 0: return [] n, m = cost_matrix.shape flat = [(cost_matrix[i, j], i, j) for i in range(n) for j in range(m)] flat.sort() row_used = set() col_used = set() pairs = [] for cost, i, j in flat: if cost >= threshold: break if i in row_used or j in col_used: continue row_used.add(i) col_used.add(j) pairs.append((i, j)) return pairs # ============================================================================= # Kalman filter box tracker (state: x, y, w, h, vx, vy, vw, vh) # ============================================================================= class KalmanBoxTracker: count = 0 def __init__(self, bbox_xyxy): KalmanBoxTracker.count += 1 self.track_id = KalmanBoxTracker.count x1, y1, x2, y2 = bbox_xyxy w, h = x2 - x1, y2 - y1 x, y = x1 + w / 2, y1 + h / 2 self.kf = _KalmanFilter() self.kf.x[:4, 0] = np.array([x, y, w, h], dtype=np.float32) self.time_since_update = 0 self.hits = 1 self.hit_streak = 1 self.age = 1 def predict(self): if self.kf.x[6] + self.kf.x[2] <= 0: self.kf.x[6] *= 0.0 self.kf.predict() self.age += 1 self.time_since_update += 1 def update(self, bbox_xyxy): self.time_since_update = 0 self.hits += 1 self.hit_streak += 1 x1, y1, x2, y2 = bbox_xyxy w, h = x2 - x1, y2 - y1 x, y = x1 + w / 2, y1 + h / 2 self.kf.update(np.array([x, y, w, h], dtype=np.float32)) def get_state(self): """Returns xyxy bbox from Kalman state.""" xx = self.kf.x[:4, 0] x, y, w, h = xx[0], xx[1], xx[2], xx[3] x1 = x - w / 2 y1 = y - h / 2 x2 = x + w / 2 y2 = y + h / 2 return np.array([x1, y1, x2, y2], dtype=np.float32) def get_cx(self): return float(self.kf.x[0, 0]) 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) for i in range(ndim): self.motion_mat[i, ndim + i] = dt self.update_mat = np.eye(ndim, 2 * ndim, dtype=np.float32) self._std_weight_position = 1.0 / 20 self._std_weight_velocity = 1.0 / 160 self.x = np.zeros((8, 1), dtype=np.float32) self.P = np.eye(8, dtype=np.float32) * 10.0 def predict(self): std_pos = [ self._std_weight_position * self.x[2], self._std_weight_position * self.x[3], self._std_weight_position * self.x[2], self._std_weight_position * self.x[3], ] std_vel = [ self._std_weight_velocity * self.x[2], self._std_weight_velocity * self.x[3], self._std_weight_velocity * self.x[2], self._std_weight_velocity * self.x[3], ] Q = np.diag(np.square(np.concatenate([std_pos, std_vel]))) self.x = self.motion_mat @ self.x 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], ] ) ) H = self.update_mat S = H @ self.P @ H.T + R K = self.P @ H.T @ np.linalg.inv(S) y = z.reshape(4, 1) - H @ self.x self.x = self.x + K @ y I_KH = np.eye(8) - K @ H self.P = I_KH @ self.P @ I_KH.T + K @ R @ K.T # ============================================================================= # 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, ): self.high_thresh = track_high_thresh self.low_thresh = track_low_thresh self.match_thresh = match_thresh self.track_buffer = track_buffer self.min_hits = min_hits self.tracked_tracks = [] self.lost_tracks = [] self.removed_tracks = [] self.frame_id = 0 def update(self, boxes_xyxy, scores): self.frame_id += 1 # --- separate detections by score --- if len(boxes_xyxy) > 0: remain = scores > self.low_thresh remain_orig_idx = np.where(remain)[0] dets = boxes_xyxy[remain] det_scores = scores[remain] is_high = det_scores > self.high_thresh is_low = ~is_high else: remain_orig_idx = np.zeros(0, dtype=np.int64) dets = np.zeros((0, 4), dtype=np.float32) det_scores = np.zeros(0, dtype=np.float32) is_high = np.zeros(0, dtype=bool) is_low = np.zeros(0, dtype=bool) # --- Kalman predict all existing tracks --- track_pool = self.tracked_tracks + self.lost_tracks num_tracks = len(track_pool) # Per-frame tracking results matched_track_idx = set() det_to_track = {} tracked_map = {} lost_map = {} # Pre-allocate these for scoping high_idx = np.array([], dtype=np.int64) low_idx = np.array([], dtype=np.int64) match_pairs_high = [] if num_tracks > 0: track_boxes = np.zeros((num_tracks, 4), dtype=np.float32) for ti, trk in enumerate(track_pool): trk.predict() track_boxes[ti] = trk.get_state() # --- first association: high-score ↔ all tracks --- high_idx = np.where(is_high)[0] high_dets = dets[is_high] unmatched_tracks = list(range(num_tracks)) if len(high_dets) > 0: iou_mat = _ious_xyxy(high_dets, track_boxes) cost_mat = 1.0 - iou_mat matches = _greedy_match(cost_mat, threshold=1.0 - self.match_thresh) for di, ti in matches: det_global = int(high_idx[di]) orig_idx = int(remain_orig_idx[det_global]) track_pool[ti].update(dets[det_global]) track_pool[ti].hit_streak = max(1, track_pool[ti].hit_streak) matched_track_idx.add(ti) det_to_track[orig_idx] = track_pool[ti].track_id 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 ] # --- second association: low-score ↔ unmatched tracks --- low_idx = np.where(is_low)[0] low_dets = dets[is_low] if len(low_dets) > 0 and len(unmatched_tracks) > 0: unmatched_boxes = track_boxes[unmatched_tracks] iou_mat = _ious_xyxy(low_dets, unmatched_boxes) cost_mat = 1.0 - iou_mat matches2 = _greedy_match( cost_mat, threshold=1.0 - self.match_thresh ) for di, uti in matches2: det_global = int(low_idx[di]) pool_idx = unmatched_tracks[uti] orig_idx = int(remain_orig_idx[det_global]) track_pool[pool_idx].update(dets[det_global]) track_pool[pool_idx].hit_streak = max( 1, track_pool[pool_idx].hit_streak ) matched_track_idx.add(pool_idx) det_to_track[orig_idx] = track_pool[pool_idx].track_id 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): if ti not in matched_track_idx: trk.hit_streak = 0 # --- lifecycle management --- new_tracked = [] new_lost = [] for trk in track_pool: if trk.time_since_update > self.track_buffer: self.removed_tracks.append(trk) elif trk.time_since_update > 0: new_lost.append(trk) else: new_tracked.append(trk) self.tracked_tracks = new_tracked self.lost_tracks = new_lost # --- confirmed tracks (both tracked and lost) --- for trk in self.tracked_tracks + self.lost_tracks: if trk.hit_streak >= self.min_hits or trk.hits >= self.min_hits: tracked_map.setdefault(trk.track_id, trk.get_cx()) for trk in self.lost_tracks: if trk.hit_streak >= self.min_hits or trk.hits >= self.min_hits: lost_map[trk.track_id] = trk.get_cx() # --- new tracks from unmatched high-score dets --- high_all = np.where(is_high)[0] matched_det_ids = set(det_to_track.keys()) for dg in high_all: orig_idx = int(remain_orig_idx[int(dg)]) if orig_idx not in matched_det_ids: trk = KalmanBoxTracker(dets[int(dg)]) self.tracked_tracks.append(trk) det_to_track[orig_idx] = trk.track_id tracked_map[trk.track_id] = trk.get_cx() return tracked_map, det_to_track, lost_map # ============================================================================= # 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, ): self.imgsz = imgsz self.conf = conf self.iou = iou self.num_classes = num_classes self.num_keypoints = num_keypoints self.score_sigmoid = score_sigmoid 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}") ret = self.rknn.init_runtime(core_mask=core_mask) if ret != 0: 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}") except Exception: pass print(f"RKNN model loaded: {model_path} imgsz={imgsz} core_mask={core_mask}") def _preprocess(self, frame): h0, w0 = frame.shape[:2] scale = min(self.imgsz / h0, self.imgsz / w0) nh, nw = int(h0 * scale), int(w0 * scale) resized = cv2.resize(frame, (nw, nh), interpolation=cv2.INTER_LINEAR) 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 rgb = cv2.cvtColor(letterbox, cv2.COLOR_BGR2RGB) gains = np.array([scale, scale, dy, dx], dtype=np.float32) return rgb, gains def __call__(self, frame): h0, w0 = frame.shape[:2] rgb, gains = self._preprocess(frame) scale, _, pad_y, pad_x = gains inp = np.expand_dims(rgb, axis=0) inp = np.ascontiguousarray(inp.astype(np.uint8)) outputs = self.rknn.inference(inputs=[inp]) if len(outputs) == 0: return [] out = outputs[0] out = np.squeeze(out, axis=0) if out.shape[0] == self.num_classes + 4: out = out.T boxes_cxcywh = out[:, :4].copy() cls_raw = out[:, 4:].copy() if self.score_sigmoid: cls_scores = 1.0 / (1.0 + np.exp(-np.clip(cls_raw, -10, 10))) 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, ) max_scores = cls_scores.max(axis=1) class_ids = cls_scores.argmax(axis=1) mask = max_scores > self.conf if mask.sum() == 0: return [] bboxes = boxes_xyxy[mask].astype(np.float32) scores = max_scores[mask].astype(np.float32) clses = class_ids[mask] bboxes[:, 0] = (bboxes[:, 0] - pad_x) / scale bboxes[:, 1] = (bboxes[:, 1] - pad_y) / scale bboxes[:, 2] = (bboxes[:, 2] - pad_x) / scale bboxes[:, 3] = (bboxes[:, 3] - pad_y) / scale bboxes[:, 0] = np.clip(bboxes[:, 0], 0, w0) bboxes[:, 1] = np.clip(bboxes[:, 1], 0, h0) bboxes[:, 2] = np.clip(bboxes[:, 2], 0, w0) bboxes[:, 3] = np.clip(bboxes[:, 3], 0, h0) detections = [] for cls_id in range(self.num_classes): idx = np.where(clses == cls_id)[0] if len(idx) == 0: continue 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, } ) return detections def release(self): self.rknn.release() # ============================================================================= # Drawing helpers # ============================================================================= def resolve_line_x(frame_width): if LINE_X is not None: return LINE_X if LINE_X_FRAC != 0.5: return int(frame_width * LINE_X_FRAC) return frame_width // 2 def crossed_line(prev_cx, cx, line_x, direction=CROSS_DIRECTION): if direction == "ltr": return prev_cx < line_x <= cx 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") def open_capture(source): 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...") for _ in range(n): cap.read() print("Stream ready!") def open_video_writer(path, w, h, fps): 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.writer = csv.writer(self.file) if new_file: self.writer.writerow(header) self.file.flush() def write_row(self, row): self.writer.writerow(row) self.file.flush() def close(self): self.file.close() class VideoSegmentWriter: def __init__(self, output_dir, w, h, fps, segment_sec): self.output_dir = Path(output_dir) self.output_dir.mkdir(parents=True, exist_ok=True) self.w, self.h, self.fps = w, h, fps self.segment_sec = segment_sec self.segment_start = time.monotonic() self.writer = None 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") def _open_next(self): if self.writer is not None: self.writer.release() 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}") def write(self, frame): if time.monotonic() - self.segment_start >= self.segment_sec: self._open_next() self.writer.write(frame) def release(self): if self.writer is not None: self.writer.release() def prune_stale_tracks(tracked, now_mono): stale = [ tid for tid, (_, ts) in tracked.items() if now_mono - ts > TRACKED_PRUNE_SEC ] for tid in stale: del tracked[tid] def overlay_rect(img, x1, y1, x2, y2, color, alpha=0.65): x1, y1 = max(0, x1), max(0, y1) x2, y2 = min(img.shape[1], x2), min(img.shape[0], y2) if x2 <= x1 or y2 <= y1: return roi = img[y1:y2, x1:x2] patch = np.full_like(roi, color, dtype=np.uint8) cv2.addWeighted(patch, alpha, roi, 1 - alpha, 0, roi) def draw_pill(img, text, x, y, bg, fg=C_TEXT, font_scale=0.45, pad_x=6, pad_y=4): font = cv2.FONT_HERSHEY_SIMPLEX (tw, th), baseline = cv2.getTextSize(text, font, font_scale, 1) x1, y1 = x, y - th - pad_y x2, y2 = x + tw + pad_x * 2, y + baseline + pad_y cv2.rectangle(img, (x1, y1), (x2, y2), bg, -1) cv2.rectangle(img, (x1, y1), (x2, y2), C_BORDER, 1) cv2.putText(img, text, (x + pad_x, y), font, font_scale, fg, 1, cv2.LINE_AA) def draw_elegant_counting_line(img, line_x, h, pulse_remaining=0): strength = pulse_remaining / max(LINE_PULSE_FRAMES, 1) glow_alpha = 0.12 + 0.18 * strength for offset in (14, 9, 5): color = tuple(int(c * glow_alpha) for c in C_LINE_GLOW) cv2.line(img, (line_x - offset, 0), (line_x - offset, h), color, 1, cv2.LINE_AA) cv2.line(img, (line_x + offset, 0), (line_x + offset, h), color, 1, cv2.LINE_AA) dash_len, gap = 18, 12 y = 0 while y < 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) 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, ) def draw_hero_count(img, line_x, h, count, pulse_remaining=0): text = str(count) font = cv2.FONT_HERSHEY_SIMPLEX boost = 0.35 * (pulse_remaining / max(COUNT_PULSE_FRAMES, 1)) font_scale, thickness = 1.4 + boost, 3 (tw, th), _ = cv2.getTextSize(text, font, font_scale, thickness) pad = 14 tx, ty = line_x - tw // 2, h // 3 + th // 2 overlay_rect( img, tx - pad, ty - th - pad, tx + tw + pad, ty + pad, C_PANEL, alpha=0.78 ) cv2.rectangle( img, (tx - pad, ty - th - pad), (tx + tw + pad, ty + pad), 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): bar_h = 40 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", (14, 14), cv2.FONT_HERSHEY_SIMPLEX, 0.32, C_MUTED, 1, cv2.LINE_AA ) batch_label = str(batch_num) if batch_num else "--" cv2.putText( img, batch_label, (14, 32), cv2.FONT_HERSHEY_SIMPLEX, 0.55, C_ACCENT, 1, cv2.LINE_AA, ) cv2.putText( img, "COUNT", (90, 14), cv2.FONT_HERSHEY_SIMPLEX, 0.32, C_MUTED, 1, cv2.LINE_AA ) cv2.putText( img, str(batch_count), (90, 32), cv2.FONT_HERSHEY_SIMPLEX, 0.55, C_GREEN, 1, cv2.LINE_AA, ) cv2.putText( img, "TOTAL", (170, 14), cv2.FONT_HERSHEY_SIMPLEX, 0.32, C_MUTED, 1, cv2.LINE_AA ) cv2.putText( img, str(total_ayam), (170, 32), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_TEXT, 1, cv2.LINE_AA, ) cv2.putText( img, "UPTIME", (250, 14), cv2.FONT_HERSHEY_SIMPLEX, 0.32, C_MUTED, 1, cv2.LINE_AA, ) cv2.putText( img, f"{elapsed_sec / 3600:.1f}h", (250, 32), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_TEXT, 1, cv2.LINE_AA, ) cv2.putText( img, "RATE", (340, 14), cv2.FONT_HERSHEY_SIMPLEX, 0.32, C_MUTED, 1, cv2.LINE_AA ) cv2.putText( img, f"{rate:.1f}/min", (340, 32), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_ACCENT, 1, cv2.LINE_AA, ) def draw_footer(img, w, h, frame_idx, live_tag, inf_ms=0.0, model_name=""): bar_h = 28 overlay_rect(img, 0, h - bar_h, w, h, C_PANEL, alpha=0.55) cv2.putText( img, f"{live_tag} | {model_name} | Frame {frame_idx} | Inf {inf_ms:.1f}ms", (12, h - 9), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA, ) def draw_skeleton_bold(img, kpts): for (a, b), color in zip(SKELETON, SK_COLORS): if a < len(kpts) and b < len(kpts): xa, ya = int(kpts[a][0]), int(kpts[a][1]) xb, yb = int(kpts[b][0]), int(kpts[b][1]) if xa > 0 and ya > 0 and xb > 0 and yb > 0: cv2.line(img, (xa, ya), (xb, yb), color, 3, cv2.LINE_AA) for kp in kpts: x, y = int(kp[0]), int(kp[1]) if x > 0 and y > 0: cv2.circle(img, (x, y), 6, (255, 255, 255), -1, cv2.LINE_AA) cv2.circle(img, (x, y), 6, (40, 40, 40), 2, cv2.LINE_AA) def draw_popups(img, popups, frame_idx): alive = [] for pop in popups: 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) 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, ) return alive def draw_batch_banner(img, w, batch_num, pulse_remaining): if pulse_remaining <= 0: return 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 ) def connect_stream(source, warmup=WARMUP_FRAMES): attempts = 0 while not shutdown_requested: cap = open_capture(source) 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...") time.sleep(RECONNECT_DELAY_SEC) continue 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)) fps = cap.get(cv2.CAP_PROP_FPS) if not fps or fps <= 1: fps = OUTPUT_FPS return cap, w, h, fps return None, 0, 0, OUTPUT_FPS # ============================================================================= # Main loop # ============================================================================= def run(): global shutdown_requested store = BatchStore( db_path=DB_PATH, state_file=STATE_FILE, camera_name=CAMERA_NAME, object_label=OBJECT_LABEL, cutoff_time=DAILY_CUTOFF_TIME, batch_timeout=BATCH_TIMEOUT_SECONDS, 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}"), ) store.start_cutoff_watcher() cross_logger = None if EXPORT_CSV: cross_logger = CsvLogger( CROSS_CSV, ["batch", "frame", "timestamp", "chicken_id"] ) model = RKNNYOLO( model_path=MODEL_PATH, core_mask=CORE_MASK, imgsz=IMGSZ, conf=CONF, num_classes=NUM_CLASSES, score_sigmoid=SCORE_SIGMOID, ) ayam_cls = int(os.getenv("AYAM_CLASS_ID", "0")) talenan_cls = int(os.getenv("TALENAN_CLASS_ID", "1")) ayam_tracker = ByteTracker( track_high_thresh=TRACK_HIGH_THRESH_0, track_low_thresh=TRACK_LOW_THRESH_0, match_thresh=TRACK_MATCH_THRESH_0, track_buffer=TRACK_BUFFER_0, min_hits=TRACK_MIN_HITS_0, ) talenan_tracker = ByteTracker( track_high_thresh=TRACK_HIGH_THRESH_1, track_low_thresh=TRACK_LOW_THRESH_1, match_thresh=TRACK_MATCH_THRESH_1, track_buffer=TRACK_BUFFER_1, min_hits=TRACK_MIN_HITS_1, ) ayam_line_crossed = set() talenan_line_crossed = set() ayam_cross_flash = {} talenan_cross_flash = {} line_pulse = count_pulse = batch_pulse = 0 popups = [] session_start = time.time() frame_idx = 0 inf_ms = 0.0 prev_gray = None video_writer = None crossing_times = deque() cap, w, h, fps = connect_stream(SOURCE) if cap is None: store.shutdown() model.release() 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 Index 0: high_thresh={TRACK_HIGH_THRESH_0} low_thresh={TRACK_LOW_THRESH_0} " f"match_thresh={TRACK_MATCH_THRESH_0} buffer={TRACK_BUFFER_0}" ) print( f"ByteTrack Index 1: high_thresh={TRACK_HIGH_THRESH_1} low_thresh={TRACK_LOW_THRESH_1} " f"match_thresh={TRACK_MATCH_THRESH_1} buffer={TRACK_BUFFER_1}" ) print(f"DB: {DB_PATH}") print(f"State: {STATE_FILE}") if RECORD_VIDEO: video_writer = VideoSegmentWriter(OUTPUT_DIR, w, h, fps, VIDEO_SEGMENT_SEC) reconnect_count = 0 while not shutdown_requested: ret, frame = cap.read() if not ret: if not IS_LIVE: break reconnect_count += 1 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) if cap is None: break line_x = resolve_line_x(w) continue now = time.time() elapsed = now - session_start mono = time.monotonic() ayam_crossed_frame = batch_closed_frame = batch_started_frame = False skip_inference = False if MOTION_DETECTION_ENABLED: gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) if prev_gray is not None: diff = cv2.absdiff(gray, prev_gray) mean_diff = cv2.mean(diff)[0] skip_inference = mean_diff < MOTION_THRESHOLD prev_gray = gray detections = [] if not skip_inference: inf_start = time.time() detections = model(frame) inf_ms = inf_ms * 0.9 + (time.time() - inf_start) * 1000 * 0.1 if detections: ayam_boxes_xyxy = [] ayam_scores = [] ayam_kpts_list = [] ayam_cx_list = [] talenan_boxes_xyxy = [] talenan_scores = [] talenan_kpts_list = [] talenan_cx_list = [] for det in detections: 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: talenan_boxes_xyxy.append(bbox) talenan_scores.append(score) talenan_kpts_list.append(kpts) talenan_cx_list.append(cx) elif cls_id == ayam_cls: ayam_boxes_xyxy.append(bbox) ayam_scores.append(score) ayam_kpts_list.append(kpts) ayam_cx_list.append(cx) ayam_boxes_xyxy = np.array(ayam_boxes_xyxy, dtype=np.float32) ayam_scores = np.array(ayam_scores, dtype=np.float32) talenan_boxes_xyxy = np.array(talenan_boxes_xyxy, 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_boxes_xyxy, ayam_scores ) talenan_track_map, talenan_det_to_track, talenan_lost_map = ( talenan_tracker.update(talenan_boxes_xyxy, talenan_scores) ) if os.getenv("DEBUG_TRACKING", "").lower() == "true": if len(ayam_boxes_xyxy) > 0 or len(talenan_boxes_xyxy) > 0: ayam_scores_str = ( f" ayam scores: {ayam_scores.round(3).tolist()}" if len(ayam_scores) > 0 else "" ) talenan_scores_str = ( f" talenan scores: {talenan_scores.round(3).tolist()}" if len(talenan_scores) > 0 else "" ) ayam_tracks_str = ( f" ayam det→track: {dict(ayam_det_to_track)}" if ayam_det_to_track else "" ) talenan_tracks_str = ( f" talenan det→track: {dict(talenan_det_to_track)}" if talenan_det_to_track else "" ) crossing_str = ( f" ayam_line_crossed: {sorted(ayam_line_crossed)}" if ayam_line_crossed else "" ) # Show cx values for one tracked object to check movement cx_sample = "" if len(ayam_cx_list) > 0 and ayam_det_to_track: sample_tid = list(ayam_det_to_track.values())[0] sample_cx = ayam_cx_list[ list(ayam_det_to_track.keys())[0] ] prev = ( ayam_tracked.get(sample_tid, (None,))[0] if ayam_tracked.get(sample_tid) else None ) cx_sample = ( f" sample tid={sample_tid} prev_cx={prev} cx={sample_cx:.1f}" ) print( f"[DEBUG F{frame_idx}] ayam_dets={len(ayam_boxes_xyxy)} " f"talenan_dets={len(talenan_boxes_xyxy)} " f"line_x={line_x}{ayam_scores_str}{talenan_scores_str}" f"{ayam_tracks_str}{talenan_tracks_str}{crossing_str}" f"{cx_sample}" ) # Process talenan crossings for di in range(len(talenan_boxes_xyxy)): tid = talenan_det_to_track.get(di) if tid is None: continue cx = talenan_cx_list[di] bbox = talenan_boxes_xyxy[di] 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 ): 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", } ) talenan_tracked[tid] = (cx, mono) # Process ayam crossings (including lost tracks for line-cross continuity) for di in range(len(ayam_boxes_xyxy)): tid = ayam_det_to_track.get(di) if tid is None: continue cx = ayam_cx_list[di] 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 os.getenv("DEBUG_TRACKING", "").lower() == "true": print( f"[DEBUG F{frame_idx}] CROSS DETECTED: tid={tid} " f"prev_cx={prev_cx:.1f} → cx={cx:.1f} line_x={line_x}" ) 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, ] ) ayam_crossed_frame = True if started_new: batch_started_frame = True crossing_times.append(mono) 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", } ) ayam_tracked[tid] = (cx, mono) # Also track lost tracks for line-crossing continuity for tid, cx in ayam_lost_map.items(): if tid not in ayam_tracked: ayam_tracked[tid] = (cx, mono) # Draw talenan for di in range(len(talenan_boxes_xyxy)): tid = talenan_det_to_track.get(di) if tid is None: continue bbox = talenan_boxes_xyxy[di] x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3]) 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 ayam for di in range(len(ayam_boxes_xyxy)): tid = ayam_det_to_track.get(di) if tid is None: continue bbox = ayam_boxes_xyxy[di] x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3]) 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) kpts = ayam_kpts_list[di] if di < len(ayam_kpts_list) else None if kpts is not None: draw_skeleton_bold(frame, kpts) if ayam_crossed_frame: line_pulse = LINE_PULSE_FRAMES count_pulse = COUNT_PULSE_FRAMES if batch_closed_frame: line_pulse = LINE_PULSE_FRAMES if batch_started_frame: batch_pulse = BATCH_PULSE_FRAMES batch_num = store.current_batch_number or 0 batch_count = store.current_batch_count display_total = store.display_total() while crossing_times and mono - crossing_times[0] > RATE_WINDOW_SEC: crossing_times.popleft() rate = (len(crossing_times) / RATE_WINDOW_SEC * 60) if crossing_times else 0.0 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, ) draw_batch_banner(frame, w, batch_num, batch_pulse) draw_footer( frame, w, h, frame_idx, "LIVE" if IS_LIVE else "FILE", inf_ms, Path(MODEL_PATH).name, ) popups = draw_popups(frame, popups, frame_idx) for flash_store in (ayam_cross_flash, talenan_cross_flash): for tid in list(flash_store): flash_store[tid] -= 1 if flash_store[tid] <= 0: del flash_store[tid] line_pulse = max(0, line_pulse - 1) count_pulse = max(0, count_pulse - 1) batch_pulse = max(0, batch_pulse - 1) if video_writer is not None: video_writer.write(frame) 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: f.write(jpeg.tobytes()) except Exception: pass 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" ) prune_stale_tracks(ayam_tracked, mono) prune_stale_tracks(talenan_tracked, mono) cap.release() if video_writer is not None: video_writer.release() if cross_logger: cross_logger.close() model.release() store.shutdown() print("\n=== Batch Summary (SQLite) ===") print(f"Database: {DB_PATH}") ayam_tracked = {} talenan_tracked = {} if __name__ == "__main__": run()