""" Edge production live counter — RTSP + YOLO RKNN + ByteTrack + zone counting. Runs on RK3588 hardware with RKNN model (320×320 input). Uses ByteTrack (Kalman filter + two-stage IoU association) for tracking. Left / right rectangular zones mark the left and right sack feeders. A sack is counted once when its centroid enters a zone. """ import numpy as np import cv2 import csv import json import os import signal import socket import threading 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 counter_store import CounterStore # --- 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_counter.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", "object") CLASS_OBJECT = os.getenv("CLASS_OBJECT", "object") # Human-readable feeder names used in count logs, e.g. # "Total karung di kandang_bawah_feeder_kanan: 5" FEEDER_LEFT_NAME = os.getenv("FEEDER_LEFT_NAME", f"{CAMERA_NAME}_feeder_kiri") FEEDER_RIGHT_NAME = os.getenv("FEEDER_RIGHT_NAME", f"{CAMERA_NAME}_feeder_kanan") DEBUG_TRACKING = os.getenv("DEBUG_TRACKING", "").lower() == "true" # Left feeder zone (fractions of frame width/height). Absolute pixel overrides # win when set (ZONE_LEFT_X1 … ZONE_LEFT_Y2). ZONE_LEFT_X1 = int(os.getenv("ZONE_LEFT_X1")) if os.getenv("ZONE_LEFT_X1") else None ZONE_LEFT_Y1 = int(os.getenv("ZONE_LEFT_Y1")) if os.getenv("ZONE_LEFT_Y1") else None ZONE_LEFT_X2 = int(os.getenv("ZONE_LEFT_X2")) if os.getenv("ZONE_LEFT_X2") else None ZONE_LEFT_Y2 = int(os.getenv("ZONE_LEFT_Y2")) if os.getenv("ZONE_LEFT_Y2") else None ZONE_LEFT_X1_FRAC = float(os.getenv("ZONE_LEFT_X1_FRAC", "0.00")) ZONE_LEFT_Y1_FRAC = float(os.getenv("ZONE_LEFT_Y1_FRAC", "0.20")) ZONE_LEFT_X2_FRAC = float(os.getenv("ZONE_LEFT_X2_FRAC", "0.35")) ZONE_LEFT_Y2_FRAC = float(os.getenv("ZONE_LEFT_Y2_FRAC", "0.85")) # Right feeder zone ZONE_RIGHT_X1 = int(os.getenv("ZONE_RIGHT_X1")) if os.getenv("ZONE_RIGHT_X1") else None ZONE_RIGHT_Y1 = int(os.getenv("ZONE_RIGHT_Y1")) if os.getenv("ZONE_RIGHT_Y1") else None ZONE_RIGHT_X2 = int(os.getenv("ZONE_RIGHT_X2")) if os.getenv("ZONE_RIGHT_X2") else None ZONE_RIGHT_Y2 = int(os.getenv("ZONE_RIGHT_Y2")) if os.getenv("ZONE_RIGHT_Y2") else None ZONE_RIGHT_X1_FRAC = float(os.getenv("ZONE_RIGHT_X1_FRAC", "0.65")) ZONE_RIGHT_Y1_FRAC = float(os.getenv("ZONE_RIGHT_Y1_FRAC", "0.20")) ZONE_RIGHT_X2_FRAC = float(os.getenv("ZONE_RIGHT_X2_FRAC", "1.00")) ZONE_RIGHT_Y2_FRAC = float(os.getenv("ZONE_RIGHT_Y2_FRAC", "0.85")) IMGSZ = int(os.getenv("IMGSZ", "320")) HALF = os.getenv("HALF", "false").lower() == "true" CONF = float(os.getenv("CONF", "0.3")) NMS_IOU = float(os.getenv("NMS_IOU", "0.45")) 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 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")) # Dedup guard against ID-switch double counts: ignore a second zone entry in the # same feeder within DEDUP_FRAMES and DEDUP_PX (centroid distance) of a recent count. DEDUP_FRAMES = int(os.getenv("DEDUP_FRAMES", "15")) DEDUP_PX = float(os.getenv("DEDUP_PX", "60")) # Trajectory inheritance across ID switches: when a new track appears, inherit the # last position of a recently-seen nearby track so zone entry is not missed when # the ID changes at the zone boundary. INHERIT_SEC = float(os.getenv("INHERIT_SEC", "1.0")) INHERIT_PX = float(os.getenv("INHERIT_PX", "60")) # Per-feeder cooldown: after a count on that feeder, ignore further counts on the # same feeder for N seconds. 0 = disabled. ZONE_COOLDOWN_LEFT_SEC = float(os.getenv("ZONE_COOLDOWN_LEFT_SEC", "0")) ZONE_COOLDOWN_RIGHT_SEC = float(os.getenv("ZONE_COOLDOWN_RIGHT_SEC", "0")) # Per-feeder dwell: sack must remain inside the zone for N seconds before it is # counted (and a cross snapshot is saved). 0 = count on first frame in zone. ZONE_DWELL_LEFT_SEC = float(os.getenv("ZONE_DWELL_LEFT_SEC", "2")) ZONE_DWELL_RIGHT_SEC = float(os.getenv("ZONE_DWELL_RIGHT_SEC", "2")) DAILY_CUTOFF_TIME = os.getenv("DAILY_CUTOFF_TIME", "20:00") EXPORT_CSV = os.getenv("EXPORT_CSV", "true").lower() == "true" CROSS_CSV = os.getenv("CROSS_CSV", f"{OUTPUT_DIR}/crossings.csv") # Save an annotated frame snapshot each time a sack finishes its zone dwell and # the counter increases. SAVE_CROSS_SNAPSHOT = os.getenv("SAVE_CROSS_SNAPSHOT", "false").lower() == "true" # Also save one snapshot the first time each object is detected (before it enters # a zone), named with the same track id so it can be correlated with the count snapshot. SAVE_DETECT_SNAPSHOT = os.getenv("SAVE_DETECT_SNAPSHOT", "false").lower() == "true" CROSS_SNAPSHOT_DIR = os.getenv("CROSS_SNAPSHOT_DIR", f"{OUTPUT_DIR}/snapshots") CROSS_SNAPSHOT_QUALITY = int(os.getenv("CROSS_SNAPSHOT_QUALITY", "85")) # Retention: delete oldest snapshots when either limit is exceeded (0 = disabled). CROSS_SNAPSHOT_MAX_FILES = int(os.getenv("CROSS_SNAPSHOT_MAX_FILES", "1000")) CROSS_SNAPSHOT_MAX_AGE_DAYS = float(os.getenv("CROSS_SNAPSHOT_MAX_AGE_DAYS", "7")) # Run the cleanup sweep at most every N seconds to limit filesystem scans. CROSS_SNAPSHOT_CLEANUP_SEC = int(os.getenv("CROSS_SNAPSHOT_CLEANUP_SEC", "3600")) 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")) 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")) # Per-pixel intensity change (0-255) for a pixel to count as "moved". MOTION_PIXEL_DELTA = int(os.getenv("MOTION_PIXEL_DELTA", "25")) # Fraction of frame pixels that must change (0-1) to trigger inference. Small, # so an object entering the edge of the frame is detected immediately. MOTION_MIN_AREA_FRAC = float(os.getenv("MOTION_MIN_AREA_FRAC", "0.002")) # Always run inference at least every N frames even if no motion (heartbeat), so a # slow/stationary object is never missed for long. MOTION_HEARTBEAT_FRAMES = int(os.getenv("MOTION_HEARTBEAT_FRAMES", "15")) # --- Runtime control (toggle counting on/off on the fly) --- # When enabled, the process watches a small JSON control file and honors its # "counting" flag. Set false to always count (ignore the control file). CONTROL_ENABLED = os.getenv("CONTROL_ENABLED", "false").lower() == "true" CONTROL_FILE = os.getenv("CONTROL_FILE", f"{OUTPUT_DIR}/control.json") # Whether counting is active on startup when no control file exists yet. CONTROL_DEFAULT_COUNTING = os.getenv("CONTROL_DEFAULT_COUNTING", "true").lower() == "true" # Re-read the control file at most every N seconds. CONTROL_POLL_SEC = float(os.getenv("CONTROL_POLL_SEC", "1.0")) # Optional TCP control socket. When enabled, the counter listens for line-based # commands so counting can be toggled over the network (in addition to the file). CONTROL_SOCKET_ENABLED = os.getenv("CONTROL_SOCKET_ENABLED", "false").lower() == "true" CONTROL_SOCKET_HOST = os.getenv("CONTROL_SOCKET_HOST", "127.0.0.1") CONTROL_SOCKET_PORT = int(os.getenv("CONTROL_SOCKET_PORT", "5090")) CROSS_FLASH_FRAMES = 12 POPUP_LIFETIME = 20 ZONE_PULSE_FRAMES = 12 COUNT_PULSE_FRAMES = 15 C_ZONE_LEFT = (80, 220, 100) # green — left feeder C_ZONE_RIGHT = (0, 165, 255) # orange — right feeder C_ZONE_FILL_ALPHA = 0.18 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_OBJECT_BOX = (0, 165, 255) 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]) def get_cy(self): return float(self.kf.x[1, 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(), track_pool[ti].get_cy()) 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(), track_pool[pool_idx].get_cy(), ) # --- 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(), trk.get_cy())) 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(), trk.get_cy()) # --- 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(), trk.get_cy()) 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 _zone_coord(abs_val, frac, span): if abs_val is not None: return int(abs_val) return int(span * frac) def resolve_zone_left(frame_width, frame_height): """Return (x1, y1, x2, y2) for the left feeder zone.""" x1 = _zone_coord(ZONE_LEFT_X1, ZONE_LEFT_X1_FRAC, frame_width) y1 = _zone_coord(ZONE_LEFT_Y1, ZONE_LEFT_Y1_FRAC, frame_height) x2 = _zone_coord(ZONE_LEFT_X2, ZONE_LEFT_X2_FRAC, frame_width) y2 = _zone_coord(ZONE_LEFT_Y2, ZONE_LEFT_Y2_FRAC, frame_height) return (min(x1, x2), min(y1, y2), max(x1, x2), max(y1, y2)) def resolve_zone_right(frame_width, frame_height): """Return (x1, y1, x2, y2) for the right feeder zone.""" x1 = _zone_coord(ZONE_RIGHT_X1, ZONE_RIGHT_X1_FRAC, frame_width) y1 = _zone_coord(ZONE_RIGHT_Y1, ZONE_RIGHT_Y1_FRAC, frame_height) x2 = _zone_coord(ZONE_RIGHT_X2, ZONE_RIGHT_X2_FRAC, frame_width) y2 = _zone_coord(ZONE_RIGHT_Y2, ZONE_RIGHT_Y2_FRAC, frame_height) return (min(x1, x2), min(y1, y2), max(x1, x2), max(y1, y2)) def point_in_zone(cx, cy, zone): x1, y1, x2, y2 = zone return x1 <= cx <= x2 and y1 <= cy <= y2 def is_duplicate_cross(recent, cx, cy, frame_idx): """True if a count near (cx, cy) happened within the dedup window (ID-switch guard).""" while recent and frame_idx - recent[0][0] > DEDUP_FRAMES: recent.popleft() for _, prev_cx, prev_cy in recent: if abs(prev_cx - cx) <= DEDUP_PX and abs(prev_cy - cy) <= DEDUP_PX: return True return False def feeder_cooldown_sec(feeder): if feeder == "left": return ZONE_COOLDOWN_LEFT_SEC return ZONE_COOLDOWN_RIGHT_SEC def feeder_dwell_sec(feeder): if feeder == "left": return ZONE_DWELL_LEFT_SEC return ZONE_DWELL_RIGHT_SEC def resolve_feeder_at(cx, cy, zone_left, zone_right): """Return 'left' | 'right' | None for the zone containing (cx, cy).""" now_left = point_in_zone(cx, cy, zone_left) now_right = point_in_zone(cx, cy, zone_right) if now_left and now_right: lx1, ly1, lx2, ly2 = zone_left rx1, ry1, rx2, ry2 = zone_right dl = abs(cx - (lx1 + lx2) / 2) + abs(cy - (ly1 + ly2) / 2) dr = abs(cx - (rx1 + rx2) / 2) + abs(cy - (ry1 + ry2) / 2) return "left" if dl <= dr else "right" if now_left: return "left" if now_right: return "right" return None def feeder_display_name(feeder): if feeder == "left": return FEEDER_LEFT_NAME return FEEDER_RIGHT_NAME def feeder_in_cooldown(feeder, mono, last_count_mono): """True if this feeder is still inside its post-count cooldown window.""" cooldown = feeder_cooldown_sec(feeder) if cooldown <= 0: return False last = last_count_mono.get(feeder) if last is None: return False return (mono - last) < cooldown def _inherit_prev(tracked, new_tid, cx, cy, mono, max_age, max_px): """Find a recently-seen nearby track for ID-switch continuation. Returns (cx, cy, mono, source_tid) or None. """ best = None best_dist = max_px for tid, (tcx, tcy, ts) in tracked.items(): if tid == new_tid: continue if mono - ts > max_age: continue if abs(tcy - cy) > max_px: continue dist = abs(tcx - cx) if dist <= best_dist: best_dist = dist best = (tcx, tcy, ts, tid) return best def _default_zone_state(): """Per-track feeder-zone state. counted: already contributed one LEFT or RIGHT count (at most one per object). side: 'left' | 'right' | None after a count. dwell_feeder / dwell_since: which zone the track is currently dwelling in, and when continuous presence in that zone began. """ return { "counted": False, "side": None, "dwell_feeder": None, "dwell_since": None, } def _copy_zone_state(src): return { "counted": bool(src.get("counted", False)), "side": src.get("side"), "dwell_feeder": src.get("dwell_feeder"), "dwell_since": src.get("dwell_since"), } def now_str(): return datetime.now().strftime("%Y-%m-%d %H:%M:%S") def read_counting_flag(default=True): """Read the 'counting' flag from the control file. Returns default on any error.""" try: with open(CONTROL_FILE, "r", encoding="utf-8") as f: data = json.load(f) return bool(data.get("counting", default)) except FileNotFoundError: return default except Exception: return default def write_control_file(counting): """Create/update the control file atomically (used to seed defaults).""" try: Path(CONTROL_FILE).parent.mkdir(parents=True, exist_ok=True) tmp = f"{CONTROL_FILE}.tmp" with open(tmp, "w", encoding="utf-8") as f: json.dump({"counting": bool(counting)}, f) os.replace(tmp, CONTROL_FILE) except Exception as exc: print(f"[{now_str()}] Failed to write control file: {exc}") def start_control_socket(): """Start a TCP server for runtime control. Commands (newline-terminated): START | RESUME | ON -> counting on STOP | PAUSE | OFF -> counting off TOGGLE -> flip STATUS | GET -> report current state It writes the shared control file, so the main loop's file-poll applies it. Returns the server socket (call .close() to stop).""" srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM) srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) srv.bind((CONTROL_SOCKET_HOST, CONTROL_SOCKET_PORT)) srv.listen(5) def handle(conn, addr): with conn: conn.settimeout(30) try: buf = b"" while not shutdown_requested: try: chunk = conn.recv(256) except socket.timeout: break if not chunk: break buf += chunk while b"\n" in buf: line, buf = buf.split(b"\n", 1) cmd = line.decode("utf-8", "ignore").strip().upper() if not cmd: continue current = read_counting_flag(CONTROL_DEFAULT_COUNTING) if cmd in ("START", "RESUME", "ON"): write_control_file(True) resp = "OK counting=on" elif cmd in ("STOP", "PAUSE", "OFF"): write_control_file(False) resp = "OK counting=off" elif cmd == "TOGGLE": write_control_file(not current) resp = f"OK counting={'off' if current else 'on'}" elif cmd in ("STATUS", "GET"): resp = f"OK counting={'on' if current else 'off'}" else: resp = "ERR unknown command" conn.sendall((resp + "\n").encode("utf-8")) except Exception: pass def loop(): print(f"Control socket listening on {CONTROL_SOCKET_HOST}:{CONTROL_SOCKET_PORT}") while not shutdown_requested: try: conn, addr = srv.accept() except OSError: break t = threading.Thread(target=handle, args=(conn, addr), daemon=True) t.start() threading.Thread(target=loop, daemon=True).start() return srv 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 cleanup_snapshots(snapshot_dir, max_files, max_age_days): """Delete oldest / expired crossing snapshots to bound disk usage.""" d = Path(snapshot_dir) if not d.is_dir(): return files = sorted(d.rglob("*.jpg"), key=lambda p: p.stat().st_mtime) if max_age_days > 0: cutoff = time.time() - max_age_days * 86400 for p in list(files): if p.stat().st_mtime < cutoff: p.unlink(missing_ok=True) files.remove(p) if max_files > 0 and len(files) > max_files: for p in files[: len(files) - max_files]: p.unlink(missing_ok=True) 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_zone(img, zone, pulse_remaining=0, label="ZONE", color=C_ZONE_LEFT): x1, y1, x2, y2 = zone strength = pulse_remaining / max(ZONE_PULSE_FRAMES, 1) fill_alpha = C_ZONE_FILL_ALPHA + 0.12 * strength overlay_rect(img, x1, y1, x2, y2, color, alpha=fill_alpha) thickness = 2 + int(2 * strength) cv2.rectangle(img, (x1, y1), (x2, y2), color, thickness, cv2.LINE_AA) # Corner ticks tick = max(12, min(x2 - x1, y2 - y1) // 8) for (ax, ay, dx, dy) in ( (x1, y1, 1, 1), (x2, y1, -1, 1), (x1, y2, 1, -1), (x2, y2, -1, -1), ): cv2.line(img, (ax, ay), (ax + dx * tick, ay), color, 2, cv2.LINE_AA) cv2.line(img, (ax, ay), (ax, ay + dy * tick), color, 2, cv2.LINE_AA) cv2.putText( img, label, (x1 + 10, y1 + 22), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1, cv2.LINE_AA, ) def draw_zone_count(img, zone, count, label, color, 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.2 + boost, 3 (tw, th), _ = cv2.getTextSize(text, font, font_scale, thickness) (lw, lh), _ = cv2.getTextSize(label, font, 0.45, 1) pad = 12 box_w = max(tw, lw) + pad * 2 box_h = th + lh + pad * 2 + 6 x1, y1, x2, y2 = zone bx1 = x1 + ((x2 - x1) - box_w) // 2 by1 = max(8, y1 - box_h - 8) bx2 = bx1 + box_w by2 = by1 + box_h overlay_rect(img, bx1, by1, bx2, by2, C_PANEL, alpha=0.78) cv2.rectangle(img, (bx1, by1), (bx2, by2), color, 2) tx = bx1 + (box_w - tw) // 2 ty = by1 + pad + th cv2.putText(img, text, (tx, ty), font, font_scale, color, thickness, cv2.LINE_AA) slx = bx1 + (box_w - lw) // 2 sly = ty + lh + 6 cv2.putText(img, label, (slx, sly), font, 0.45, C_MUTED, 1, cv2.LINE_AA) def draw_hud(img, w, total, total_left, total_right, 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, "TOTAL", (14, 14), cv2.FONT_HERSHEY_SIMPLEX, 0.32, C_MUTED, 1, cv2.LINE_AA ) cv2.putText( img, str(total), (14, 32), cv2.FONT_HERSHEY_SIMPLEX, 0.55, C_GREEN, 1, cv2.LINE_AA, ) cv2.putText( img, "LEFT", (100, 14), cv2.FONT_HERSHEY_SIMPLEX, 0.32, C_MUTED, 1, cv2.LINE_AA ) cv2.putText( img, str(total_left), (100, 32), cv2.FONT_HERSHEY_SIMPLEX, 0.55, C_ZONE_LEFT, 1, cv2.LINE_AA, ) cv2.putText( img, "RIGHT", (180, 14), cv2.FONT_HERSHEY_SIMPLEX, 0.32, C_MUTED, 1, cv2.LINE_AA ) cv2.putText( img, str(total_right), (180, 32), cv2.FONT_HERSHEY_SIMPLEX, 0.55, C_ZONE_RIGHT, 1, cv2.LINE_AA, ) cv2.putText( img, "UPTIME", (270, 14), cv2.FONT_HERSHEY_SIMPLEX, 0.32, C_MUTED, 1, cv2.LINE_AA, ) cv2.putText( img, f"{elapsed_sec / 3600:.1f}h", (270, 32), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_TEXT, 1, cv2.LINE_AA, ) cv2.putText( img, "RATE", (360, 14), cv2.FONT_HERSHEY_SIMPLEX, 0.32, C_MUTED, 1, cv2.LINE_AA ) cv2.putText( img, f"{rate:.1f}/min", (360, 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 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 = CounterStore( db_path=DB_PATH, state_file=STATE_FILE, camera_name=CAMERA_NAME, object_label=OBJECT_LABEL, cutoff_time=DAILY_CUTOFF_TIME, logger=lambda msg: print(f"[{now_str()}] {msg}"), ) store.start_cutoff_watcher() cross_logger = None if EXPORT_CSV: cross_logger = CsvLogger( CROSS_CSV, ["counting_date", "frame", "feeder", "object_id"] ) model = RKNNYOLO( model_path=MODEL_PATH, core_mask=CORE_MASK, imgsz=IMGSZ, conf=CONF, iou=NMS_IOU, num_classes=NUM_CLASSES, score_sigmoid=SCORE_SIGMOID, ) object_cls = int(os.getenv("OBJECT_CLASS_ID", "0")) object_tracker = ByteTracker( track_high_thresh=TRACK_HIGH_THRESH, track_low_thresh=TRACK_LOW_THRESH, match_thresh=TRACK_MATCH_THRESH, track_buffer=TRACK_BUFFER, min_hits=TRACK_MIN_HITS, ) # Per-track zone state: enter left → LEFT feeder; enter right → RIGHT feeder. # Keys survive ID switches via inheritance (see _inherit_prev). object_zone_state = {} recent_cross_left = deque() recent_cross_right = deque() # Monotonic timestamp of last successful count per feeder (for cooldown). last_count_mono = {"left": None, "right": None} detect_snapshot_ids = set() object_cross_flash = {} zone_pulse = count_left_pulse = count_right_pulse = 0 popups = [] session_start = time.time() frame_idx = 0 inf_ms = 0.0 prev_gray = None frames_since_infer = 0 video_writer = None crossing_times = deque() counter_left = 0 counter_right = 0 last_snapshot_cleanup = 0.0 counting_active = True last_control_poll = 0.0 control_socket = None if CONTROL_ENABLED: if not Path(CONTROL_FILE).exists(): write_control_file(CONTROL_DEFAULT_COUNTING) counting_active = read_counting_flag(CONTROL_DEFAULT_COUNTING) print( f"Runtime control enabled | file={CONTROL_FILE} | " f"counting={'ON' if counting_active else 'OFF'}" ) if CONTROL_SOCKET_ENABLED: try: control_socket = start_control_socket() except Exception as exc: print(f"[{now_str()}] Failed to start control socket: {exc}") cap, w, h, fps = connect_stream(SOURCE) if cap is None: store.shutdown() model.release() return zone_left = resolve_zone_left(w, h) zone_right = resolve_zone_right(w, h) print( f"RKNN+ByteTrack zone counter | {w}x{h} @ {fps}fps | " f"LEFT feeder {zone_left} RIGHT feeder {zone_right}" ) print( f"Cooldownoldown: left={ZONE_COOLDOWN_LEFT_SEC}s right={ZONE_COOLDOWN_RIGHT_SEC}s " f"(0=off) | Dwell: left={ZONE_DWELL_LEFT_SEC}s right={ZONE_DWELL_RIGHT_SEC}s " f"(0=immediate)" ) 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) 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 zone_left = resolve_zone_left(w, h) zone_right = resolve_zone_right(w, h) continue now = time.time() elapsed = now - session_start mono = time.monotonic() object_crossed_frame = False cross_events_frame = [] detect_events_frame = [] if CONTROL_ENABLED and (now - last_control_poll) >= CONTROL_POLL_SEC: last_control_poll = now new_flag = read_counting_flag(CONTROL_DEFAULT_COUNTING) if new_flag != counting_active: counting_active = new_flag print(f"[{now_str()}] Counting {'RESUMED' if counting_active else 'PAUSED'} via control file") skip_inference = False if not counting_active: skip_inference = True elif MOTION_DETECTION_ENABLED: gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) if prev_gray is not None: diff = cv2.absdiff(gray, prev_gray) moved = int(np.count_nonzero(diff > MOTION_PIXEL_DELTA)) moved_frac = moved / diff.size skip_inference = moved_frac < MOTION_MIN_AREA_FRAC if frames_since_infer >= MOTION_HEARTBEAT_FRAMES: skip_inference = False prev_gray = gray detections = [] if not skip_inference: frames_since_infer = 0 inf_start = time.time() detections = model(frame) inf_ms = inf_ms * 0.9 + (time.time() - inf_start) * 1000 * 0.1 else: frames_since_infer += 1 object_boxes_xyxy = [] object_scores = [] object_kpts_list = [] object_cx_list = [] object_cy_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 cy = (bbox[1] + bbox[3]) / 2.0 if cls_id == object_cls: object_boxes_xyxy.append(bbox) object_scores.append(score) object_kpts_list.append(kpts) object_cx_list.append(cx) object_cy_list.append(cy) object_boxes_xyxy = np.array(object_boxes_xyxy, dtype=np.float32).reshape(-1, 4) object_scores = np.array(object_scores, dtype=np.float32) object_track_map, object_det_to_track, object_lost_map = object_tracker.update( object_boxes_xyxy, object_scores ) if DEBUG_TRACKING and len(object_boxes_xyxy) > 0: scores_str = ( f" scores: {object_scores.round(3).tolist()}" if len(object_scores) > 0 else "" ) tracks_str = ( f" det->track: {dict(object_det_to_track)}" if object_det_to_track else "" ) zone_str = "" if object_zone_state: left_ids = sorted( tid for tid, st in object_zone_state.items() if st.get("side") == "left" ) right_ids = sorted( tid for tid, st in object_zone_state.items() if st.get("side") == "right" ) if left_ids: zone_str += f" left_counted: {left_ids}" if right_ids: zone_str += f" right_counted: {right_ids}" print( f"[DEBUG F{frame_idx}] dets={len(object_boxes_xyxy)} " f"tracks={len(object_track_map)} " f"zone_L={zone_left} zone_R={zone_right}{scores_str}" f"{tracks_str}{zone_str}", flush=True, ) # --- zone dwell counting only on DETECTED objects this frame --- # A sack must remain inside a feeder zone for ZONE_DWELL_*_SEC before # it is counted (and a cross snapshot is saved). Leaving the zone # resets the dwell timer. Cooldown still applies after a successful count. for di in range(len(object_boxes_xyxy)): tid = object_det_to_track.get(di) if tid is None: continue cx = object_cx_list[di] cy = object_cy_list[di] if tid not in detect_snapshot_ids: detect_snapshot_ids.add(tid) detect_events_frame.append(tid) if tid not in object_tracked: inherited = _inherit_prev( object_tracked, tid, cx, cy, mono, INHERIT_SEC, INHERIT_PX ) if inherited is not None: object_tracked[tid] = inherited[:3] src_tid = inherited[3] if src_tid in object_zone_state: object_zone_state[tid] = _copy_zone_state( object_zone_state[src_tid] ) if DEBUG_TRACKING: src_st = object_zone_state.get(tid, {}) print( f"[DEBUG F{frame_idx}] INHERIT prev for new tid={tid} " f"from tid={src_tid} ({inherited[0]:.1f},{inherited[1]:.1f})" f" counted={src_st.get('counted', False)}" f" side={src_st.get('side')}" f" dwell={src_st.get('dwell_feeder')}", flush=True, ) if tid not in object_zone_state: object_zone_state[tid] = _default_zone_state() st = object_zone_state[tid] feeder = resolve_feeder_at(cx, cy, zone_left, zone_right) if st["counted"]: object_tracked[tid] = (cx, cy, mono) continue if feeder is None: # Left all zones — reset dwell progress. if st["dwell_feeder"] is not None and DEBUG_TRACKING: print( f"[DEBUG F{frame_idx}] DWELL RESET: tid={tid} " f"left feeder={st['dwell_feeder']}", flush=True, ) st["dwell_feeder"] = None st["dwell_since"] = None object_tracked[tid] = (cx, cy, mono) continue if st["dwell_feeder"] != feeder: st["dwell_feeder"] = feeder st["dwell_since"] = mono if DEBUG_TRACKING: print( f"[DEBUG F{frame_idx}] DWELL START: tid={tid} " f"feeder={feeder} need={feeder_dwell_sec(feeder):.1f}s", flush=True, ) object_tracked[tid] = (cx, cy, mono) continue dwell_need = feeder_dwell_sec(feeder) dwell_elapsed = mono - (st["dwell_since"] or mono) if dwell_need > 0 and dwell_elapsed < dwell_need: object_tracked[tid] = (cx, cy, mono) continue # Dwell complete — apply dedup / cooldown, then count + snapshot. recent = recent_cross_left if feeder == "left" else recent_cross_right if is_duplicate_cross(recent, cx, cy, frame_idx): st["counted"] = True st["side"] = feeder if DEBUG_TRACKING: print( f"[DEBUG F{frame_idx}] DUP ZONE IGNORED: tid={tid} " f"cx={cx:.1f} cy={cy:.1f} feeder={feeder}", flush=True, ) object_tracked[tid] = (cx, cy, mono) continue if feeder_in_cooldown(feeder, mono, last_count_mono): # Keep dwelling; count once cooldown expires if still in zone. if DEBUG_TRACKING: cool = feeder_cooldown_sec(feeder) elapsed = mono - (last_count_mono.get(feeder) or mono) print( f"[DEBUG F{frame_idx}] DWELL WAIT COOLDOWN: tid={tid} " f"feeder={feeder} elapsed={elapsed:.2f}s / {cool}s", flush=True, ) object_tracked[tid] = (cx, cy, mono) continue if DEBUG_TRACKING: print( f"[DEBUG F{frame_idx}] DWELL COUNT: tid={tid} " f"cx={cx:.1f} cy={cy:.1f} feeder={feeder} " f"dwell={dwell_elapsed:.2f}s", flush=True, ) recent.append((frame_idx, cx, cy)) st["counted"] = True st["side"] = feeder last_count_mono[feeder] = mono if feeder == "left": counter_left += 1 count_left_pulse = COUNT_PULSE_FRAMES feeder_total = counter_left else: counter_right += 1 count_right_pulse = COUNT_PULSE_FRAMES feeder_total = counter_right # Frigate-style count log, e.g. # "Total karung di kandang_bawah_feeder_kanan: 5" print( f"Total {OBJECT_LABEL} di {feeder_display_name(feeder)}: {feeder_total}", flush=True, ) store.record_zone_entry(tid, feeder) if cross_logger: cross_logger.write_row( [ store.get_counting_date(), frame_idx, feeder, tid, ] ) object_crossed_frame = True cross_events_frame.append((tid, feeder)) crossing_times.append(mono) object_cross_flash[tid] = CROSS_FLASH_FRAMES popups.append( { "x": int(cx) - 12, "y": int(cy), "born": frame_idx, "text": f"+1 {feeder.upper()}", } ) object_tracked[tid] = (cx, cy, mono) for tid, (cx, cy) in object_lost_map.items(): if tid not in object_tracked: object_tracked[tid] = (cx, cy, mono) for di in range(len(object_boxes_xyxy)): tid = object_det_to_track.get(di) if tid is None: continue bbox = object_boxes_xyxy[di] x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3]) flash = object_cross_flash.get(tid, 0) color = C_GREEN if flash > 0 else C_OBJECT_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 = object_kpts_list[di] if di < len(object_kpts_list) else None if kpts is not None: draw_skeleton_bold(frame, kpts) if object_crossed_frame: zone_pulse = ZONE_PULSE_FRAMES 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_zone( frame, zone_left, zone_pulse, label="LEFT FEEDER", color=C_ZONE_LEFT ) draw_elegant_zone( frame, zone_right, zone_pulse, label="RIGHT FEEDER", color=C_ZONE_RIGHT ) draw_zone_count( frame, zone_left, counter_left, "LEFT", C_ZONE_LEFT, count_left_pulse ) draw_zone_count( frame, zone_right, counter_right, "RIGHT", C_ZONE_RIGHT, count_right_pulse ) draw_hud( frame, w, counter_left + counter_right, counter_left, counter_right, elapsed, rate, ) 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) if CONTROL_ENABLED and not counting_active: badge = "COUNTING PAUSED" (bw, bh), _ = cv2.getTextSize(badge, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2) bx = w // 2 - bw // 2 overlay_rect(frame, bx - 14, 48, bx + bw + 14, 48 + bh + 18, C_PANEL, alpha=0.75) cv2.rectangle(frame, (bx - 14, 48), (bx + bw + 14, 48 + bh + 18), C_ACCENT, 2) cv2.putText(frame, badge, (bx, 48 + bh + 6), cv2.FONT_HERSHEY_SIMPLEX, 0.6, C_ACCENT, 2, cv2.LINE_AA) for tid in list(object_cross_flash): object_cross_flash[tid] -= 1 if object_cross_flash[tid] <= 0: del object_cross_flash[tid] zone_pulse = max(0, zone_pulse - 1) count_left_pulse = max(0, count_left_pulse - 1) count_right_pulse = max(0, count_right_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) ok, jpeg = cv2.imencode( ".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, LIVE_STREAM_QUALITY] ) if ok: tmp_path = f"{LIVE_STREAM_FRAME_PATH}.tmp" with open(tmp_path, "wb") as f: f.write(jpeg.tobytes()) os.replace(tmp_path, LIVE_STREAM_FRAME_PATH) except Exception: pass if (SAVE_DETECT_SNAPSHOT and detect_events_frame) or ( SAVE_CROSS_SNAPSHOT and cross_events_frame ): try: ts = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3] if SAVE_DETECT_SNAPSHOT and detect_events_frame: detect_dir = Path(CROSS_SNAPSHOT_DIR) / "detect" detect_dir.mkdir(parents=True, exist_ok=True) for tid in detect_events_frame: fname = f"{ts}_detect_id{tid}_f{frame_idx}.jpg" cv2.imwrite( str(detect_dir / fname), frame, [cv2.IMWRITE_JPEG_QUALITY, CROSS_SNAPSHOT_QUALITY], ) if SAVE_CROSS_SNAPSHOT and cross_events_frame: cross_dir = Path(CROSS_SNAPSHOT_DIR) / "cross" cross_dir.mkdir(parents=True, exist_ok=True) for tid, feeder in cross_events_frame: fname = f"{ts}_{feeder}_id{tid}_f{frame_idx}.jpg" cv2.imwrite( str(cross_dir / fname), frame, [cv2.IMWRITE_JPEG_QUALITY, CROSS_SNAPSHOT_QUALITY], ) if now - last_snapshot_cleanup >= CROSS_SNAPSHOT_CLEANUP_SEC: cleanup_snapshots( CROSS_SNAPSHOT_DIR, CROSS_SNAPSHOT_MAX_FILES, CROSS_SNAPSHOT_MAX_AGE_DAYS, ) last_snapshot_cleanup = now except Exception as exc: print(f"[{now_str()}] Failed to save snapshot: {exc}") frame_idx += 1 prune_stale_tracks(object_tracked, mono) cap.release() if video_writer is not None: video_writer.release() if cross_logger: cross_logger.close() if control_socket is not None: try: control_socket.close() except Exception: pass model.release() store.shutdown() print("\n=== Daily Counter Summary (SQLite) ===") print(f"Database: {DB_PATH}") object_tracked = {} if __name__ == "__main__": run()