1789 lines
63 KiB
Python
1789 lines
63 KiB
Python
"""
|
||
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() |