""" Edge production live counter — RTSP + YOLO RKNN + line crossing. Runs on RK3588 hardware with RKNN model (320×320 input). Replaces the Jetson/TensorRT variant. """ import numpy as np import cv2 import csv import os import signal import time from datetime import datetime from pathlib import Path from dotenv import load_dotenv load_dotenv() from rknnlite.api import RKNNLite from batch_store import BatchStore # --- config (override via env / .env) --- OUTPUT_DIR = os.getenv('OUTPUT_DIR', '/opt/jetson-counter') DB_PATH = os.getenv('DB_PATH', f'{OUTPUT_DIR}/jetson_counter.db') STATE_FILE = os.getenv('STATE_FILE', f'{OUTPUT_DIR}/current_batch.json') SOURCE = os.getenv('SOURCE', 'rtsp://user:pass@192.168.0.100:554/stream1') MODEL_PATH = os.getenv('MODEL_PATH', '/opt/jetson-counter/yolo11n.rknn') CAMERA_NAME = os.getenv('CAMERA_NAME', 'CC1') OBJECT_LABEL = os.getenv('OBJECT_LABEL', 'ayam-potong') CLASS_AYAM = os.getenv('CLASS_AYAM', 'ayam') CLASS_TALENAN = os.getenv('CLASS_TALENAN', 'talenan') LINE_X = int(os.getenv('LINE_X')) if os.getenv('LINE_X') else None LINE_X_FRAC = float(os.getenv('LINE_X_FRAC', '0.5')) CROSS_DIRECTION = os.getenv('CROSS_DIRECTION', 'rtl').lower() IMGSZ = int(os.getenv('IMGSZ', '320')) HALF = os.getenv('HALF', 'false').lower() == 'true' CONF = float(os.getenv('CONF', '0.3')) DEVICE = int(os.getenv('DEVICE', '0')) # RKNN-specific — core mask for NPU # 1 = core0, 2 = core1, 3 = core0+core1 (dual), 7 = all three 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' DAILY_CUTOFF_TIME = os.getenv('DAILY_CUTOFF_TIME', '20:00') BATCH_TIMEOUT_SECONDS = float(os.getenv('BATCH_TIMEOUT_SECONDS', '300')) IGNORE_BATCH_LABEL_TIMEOUT = float(os.getenv('IGNORE_BATCH_LABEL_TIMEOUT_SECONDS', '30')) MIN_OBJECT_PER_BATCH = int(os.getenv('MIN_OBJECT_PER_BATCH', '60')) MIN_DURATION_PER_BATCH = int(os.getenv('MIN_DURATION_PER_BATCH', '60')) EXPORT_CSV = os.getenv('EXPORT_CSV', 'true').lower() == 'true' CROSS_CSV = os.getenv('CROSS_CSV', f'{OUTPUT_DIR}/batch_crossings.csv') WARMUP_FRAMES = int(os.getenv('WARMUP_FRAMES', '30')) RECONNECT_DELAY_SEC = int(os.getenv('RECONNECT_DELAY_SEC', '3')) MAX_RECONNECT_ATTEMPTS = int(os.getenv('MAX_RECONNECT_ATTEMPTS', '0')) FLUSH_EVERY_N_FRAMES = int(os.getenv('FLUSH_EVERY_N_FRAMES', '100')) TRACKED_PRUNE_SEC = int(os.getenv('TRACKED_PRUNE_SEC', '300')) RECORD_VIDEO = os.getenv('RECORD_VIDEO', 'false').lower() == 'true' VIDEO_SEGMENT_SEC = int(os.getenv('VIDEO_SEGMENT_SEC', '3600')) OUTPUT_FPS = int(os.getenv('OUTPUT_FPS', '15')) LIVE_STREAM_ENABLED = os.getenv('LIVE_STREAM_ENABLED', 'false').lower() == 'true' LIVE_STREAM_FRAME_PATH = os.getenv('LIVE_STREAM_FRAME_PATH', '/dev/shm/jetson-counter/live_frame.jpg') LIVE_STREAM_QUALITY = int(os.getenv('LIVE_STREAM_QUALITY', '75')) LIVE_STREAM_EVERY_N = int(os.getenv('LIVE_STREAM_EVERY_N', '2')) RTSP_FFMPEG_OPTIONS = os.getenv( 'OPENCV_FFMPEG_CAPTURE_OPTIONS', 'rtsp_transport;tcp|fflags;nobuffer|flags;low_delay', ) IS_LIVE = SOURCE.lower().startswith(('rtsp://', 'http://')) CROSS_FLASH_FRAMES = 12 POPUP_LIFETIME = 20 LINE_PULSE_FRAMES = 12 COUNT_PULSE_FRAMES = 15 BATCH_PULSE_FRAMES = 20 SKELETON = [(0, 1), (4, 3), (1, 2), (3, 2), (2, 6), (2, 5), (2, 7), (7, 8)] SK_COLORS = [ (0, 255, 255), (0, 255, 255), (255, 0, 255), (255, 0, 255), (0, 255, 0), (255, 255, 0), (0, 0, 255), (200, 200, 0), ] C_PANEL = (28, 24, 18) C_BORDER = (90, 85, 75) C_ACCENT = (255, 200, 60) C_GREEN = (80, 220, 100) C_TEXT = (235, 235, 235) C_MUTED = (150, 150, 150) C_AYAM_BOX = (0, 165, 255) C_TALENAN_BOX = (220, 120, 60) C_LINE_CORE = (180, 220, 255) C_LINE_GLOW = (100, 160, 220) shutdown_requested = False def request_shutdown(signum, frame): global shutdown_requested shutdown_requested = True print('\nShutdown requested — finishing current frame...') signal.signal(signal.SIGINT, request_shutdown) signal.signal(signal.SIGTERM, request_shutdown) # ============================================================================= # YOLO output decoder (NMS only — boxes are pre-decoded by the model) # ============================================================================= def _nms(boxes, scores, iou_thr=0.45): order = np.argsort(scores)[::-1] keep = [] while len(order) > 0: idx = order[0] keep.append(idx) if len(order) == 1: break xx1 = np.maximum(boxes[idx, 0], boxes[order[1:], 0]) yy1 = np.maximum(boxes[idx, 1], boxes[order[1:], 1]) xx2 = np.minimum(boxes[idx, 2], boxes[order[1:], 2]) yy2 = np.minimum(boxes[idx, 3], boxes[order[1:], 3]) w = np.maximum(0.0, xx2 - xx1) h = np.maximum(0.0, yy2 - yy1) inter = w * h area_i = (boxes[idx, 2] - boxes[idx, 0]) * (boxes[idx, 3] - boxes[idx, 1]) area_o = (boxes[order[1:], 2] - boxes[order[1:], 0]) * (boxes[order[1:], 3] - boxes[order[1:], 1]) iou = inter / (area_i + area_o - inter + 1e-16) order = order[1:][iou < iou_thr] return np.array(keep) def _compute_iou(box1, boxes2): """IoU of one box (cxcywh) against a set (2D) or single box (1D).""" if boxes2.ndim == 1: boxes2 = boxes2.reshape(1, -1) cx, cy, w, h = box1 x1, y1 = cx - w / 2, cy - h / 2 x2, y2 = cx + w / 2, cy + h / 2 area1 = w * h cxs, cys, ws, hs = boxes2[:, 0], boxes2[:, 1], boxes2[:, 2], boxes2[:, 3] x1s, y1s = cxs - ws / 2, cys - hs / 2 x2s, y2s = cxs + ws / 2, cys + hs / 2 areas2 = ws * hs xx1 = np.maximum(x1, x1s) yy1 = np.maximum(y1, y1s) xx2 = np.minimum(x2, x2s) yy2 = np.minimum(y2, y2s) inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1) return inter / (area1 + areas2 - inter + 1e-16) # ============================================================================= # Simple IoU tracker (replaces bytetrack — same persist behaviour) # ============================================================================= class SimpleTracker: def __init__(self, max_age=30, min_hits=1, iou_threshold=0.3): self.max_age = max_age self.min_hits = min_hits self.iou_threshold = iou_threshold self.tracks = {} # track_id -> {box, cx, age, hits, time_since_update} self.next_id = 1 def update(self, detections): """detections: list of (cx, box_cxcywh). Returns (track_map, det_to_track).""" now = time.monotonic() for tid in self.tracks: self.tracks[tid]['time_since_update'] += 1 matched_det = set() matched_track = set() assignments = [] # (track_id, det_idx) det_to_track = {} # det_idx → track_id if detections and self.tracks: track_ids = list(self.tracks.keys()) track_boxes = np.stack([self.tracks[t]['box'] for t in track_ids], axis=0) for di, det in enumerate(detections): _, det_box = det ious = np.array([_compute_iou(det_box, track_boxes[t:t + 1]) for t in range(len(track_ids))]) best_j = int(np.argmax(ious)) if ious[best_j] >= self.iou_threshold and track_ids[best_j] not in matched_track: assignments.append((track_ids[best_j], di)) matched_track.add(track_ids[best_j]) matched_det.add(di) for tid, di in assignments: cx, box = detections[di] self.tracks[tid]['cx'] = cx self.tracks[tid]['box'] = box self.tracks[tid]['hits'] += 1 self.tracks[tid]['time_since_update'] = 0 self.tracks[tid]['last_update'] = now det_to_track[di] = tid for di, det in enumerate(detections): if di not in matched_det: cx, box = det new_id = self.next_id self.next_id += 1 self.tracks[new_id] = { 'cx': cx, 'box': box, 'hits': 1, 'time_since_update': 0, 'last_update': now, } det_to_track[di] = new_id stale = [tid for tid, t in self.tracks.items() if t['time_since_update'] > self.max_age] for tid in stale: del self.tracks[tid] track_map = {tid: self.tracks[tid]['cx'] for tid in self.tracks if self.tracks[tid]['hits'] >= self.min_hits} return track_map, det_to_track # ============================================================================= # 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: from rknnlite.api import RKNNLite as _RK 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): """Letterbox-resize to imgsz×imgsz, maintain aspect ratio, BGR→RGB, normalize.""" 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): """Run inference on BGR frame. Returns list of detection dicts.""" 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] # (1, 4+num_classes, N) or (1, N, 4+num_classes) out = np.squeeze(out, axis=0) # (C, N) or (N, C) if out.shape[0] == self.num_classes + 4: out = out.T # (C, N) → (N, C) boxes_cxcywh = out[:, :4].copy() # cx, cy, w, h at model resolution 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 (unchanged from original) # ============================================================================= def resolve_line_x(frame_width): if LINE_X is not None: return LINE_X if LINE_X_FRAC != 0.5: return int(frame_width * LINE_X_FRAC) return frame_width // 2 def crossed_line(prev_cx, cx, line_x, direction=CROSS_DIRECTION): if direction == 'ltr': return prev_cx < line_x <= cx if direction == 'both': return (prev_cx > line_x >= cx) or (prev_cx < line_x <= cx) return prev_cx > line_x >= cx def now_str(): return datetime.now().strftime('%Y-%m-%d %H:%M:%S') def open_capture(source): if source.lower().startswith(('rtsp://', 'http://')): os.environ['OPENCV_FFMPEG_CAPTURE_OPTIONS'] = RTSP_FFMPEG_OPTIONS cap = cv2.VideoCapture(source, cv2.CAP_FFMPEG) cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) return cap def warmup_stream(cap, n=WARMUP_FRAMES): print('Warming up stream...') for _ in range(n): cap.read() print('Stream ready!') def open_video_writer(path, w, h, fps): return cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*'avc1'), fps, (w, h)) class CsvLogger: def __init__(self, path, header): Path(path).parent.mkdir(parents=True, exist_ok=True) new_file = not Path(path).exists() or Path(path).stat().st_size == 0 self.file = open(path, 'a', newline='', buffering=1) self.writer = csv.writer(self.file) if new_file: self.writer.writerow(header) self.file.flush() def write_row(self, row): self.writer.writerow(row) self.file.flush() def close(self): self.file.close() class VideoSegmentWriter: def __init__(self, output_dir, w, h, fps, segment_sec): self.output_dir = Path(output_dir) self.output_dir.mkdir(parents=True, exist_ok=True) self.w, self.h, self.fps = w, h, fps self.segment_sec = segment_sec self.segment_start = time.monotonic() self.writer = None self._open_next() def _segment_path(self): ts = datetime.now().strftime('%Y%m%d_%H%M%S') return str(self.output_dir / f'live_{ts}.mp4') def _open_next(self): if self.writer is not None: self.writer.release() path = self._segment_path() self.writer = open_video_writer(path, self.w, self.h, self.fps) self.segment_start = time.monotonic() print(f'Recording segment: {path}') def write(self, frame): if time.monotonic() - self.segment_start >= self.segment_sec: self._open_next() self.writer.write(frame) def release(self): if self.writer is not None: self.writer.release() def prune_stale_tracks(tracked, now_mono): stale = [tid for tid, (_, ts) in tracked.items() if now_mono - ts > TRACKED_PRUNE_SEC] for tid in stale: del tracked[tid] def overlay_rect(img, x1, y1, x2, y2, color, alpha=0.65): x1, y1 = max(0, x1), max(0, y1) x2, y2 = min(img.shape[1], x2), min(img.shape[0], y2) if x2 <= x1 or y2 <= y1: return roi = img[y1:y2, x1:x2] patch = np.full_like(roi, color, dtype=np.uint8) cv2.addWeighted(patch, alpha, roi, 1 - alpha, 0, roi) def draw_pill(img, text, x, y, bg, fg=C_TEXT, font_scale=0.45, pad_x=6, pad_y=4): font = cv2.FONT_HERSHEY_SIMPLEX (tw, th), baseline = cv2.getTextSize(text, font, font_scale, 1) x1, y1 = x, y - th - pad_y x2, y2 = x + tw + pad_x * 2, y + baseline + pad_y cv2.rectangle(img, (x1, y1), (x2, y2), bg, -1) cv2.rectangle(img, (x1, y1), (x2, y2), C_BORDER, 1) cv2.putText(img, text, (x + pad_x, y), font, font_scale, fg, 1, cv2.LINE_AA) def draw_elegant_counting_line(img, line_x, h, pulse_remaining=0): strength = pulse_remaining / max(LINE_PULSE_FRAMES, 1) glow_alpha = 0.12 + 0.18 * strength for offset in (14, 9, 5): color = tuple(int(c * glow_alpha) for c in C_LINE_GLOW) cv2.line(img, (line_x - offset, 0), (line_x - offset, h), color, 1, cv2.LINE_AA) cv2.line(img, (line_x + offset, 0), (line_x + offset, h), color, 1, cv2.LINE_AA) dash_len, gap = 18, 12 y = 0 while y < h: y_end = min(y + dash_len, h) cv2.line(img, (line_x, y), (line_x, y_end), C_LINE_CORE, 2, cv2.LINE_AA) y += dash_len + gap cv2.putText(img, 'COUNT LINE', (line_x - 46, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.42, C_LINE_CORE, 1, cv2.LINE_AA) def draw_hero_count(img, line_x, h, count, pulse_remaining=0): text = str(count) font = cv2.FONT_HERSHEY_SIMPLEX boost = 0.35 * (pulse_remaining / max(COUNT_PULSE_FRAMES, 1)) font_scale, thickness = 1.6 + boost, 3 (tw, th), _ = cv2.getTextSize(text, font, font_scale, thickness) pad = 14 tx, ty = line_x - tw // 2, h // 2 + th // 2 overlay_rect(img, tx - pad, ty - th - pad, tx + tw + pad, ty + pad // 2, C_PANEL, alpha=0.78) cv2.rectangle(img, (tx - pad, ty - th - pad), (tx + tw + pad, ty + pad // 2), C_LINE_CORE, 2) cv2.putText(img, text, (tx, ty), font, font_scale, C_GREEN, thickness, cv2.LINE_AA) def draw_hud(img, w, batch_num, batch_count, total_ayam, elapsed_sec, rate, camera_id, clock): bar_h = 52 overlay_rect(img, 0, 0, w, bar_h, C_PANEL, alpha=0.72) cv2.line(img, (0, bar_h), (w, bar_h), C_BORDER, 1) cv2.putText(img, 'BATCH', (16, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) batch_label = str(batch_num) if batch_num else '—' cv2.putText(img, batch_label, (16, 44), cv2.FONT_HERSHEY_SIMPLEX, 0.9, C_ACCENT, 2, cv2.LINE_AA) cv2.putText(img, 'COUNT', (100, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) cv2.putText(img, str(batch_count), (100, 44), cv2.FONT_HERSHEY_SIMPLEX, 0.9, C_GREEN, 2, cv2.LINE_AA) cv2.putText(img, 'TOTAL', (190, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) cv2.putText(img, str(total_ayam), (190, 44), cv2.FONT_HERSHEY_SIMPLEX, 0.7, C_TEXT, 1, cv2.LINE_AA) cv2.putText(img, 'UPTIME', (280, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) cv2.putText(img, f'{elapsed_sec / 3600:.1f}h', (280, 44), cv2.FONT_HERSHEY_SIMPLEX, 0.7, C_TEXT, 1, cv2.LINE_AA) cv2.putText(img, 'RATE', (380, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) cv2.putText(img, f'{rate:.1f}/min', (380, 44), cv2.FONT_HERSHEY_SIMPLEX, 0.7, C_ACCENT, 1, cv2.LINE_AA) cv2.putText(img, clock, (w - 180, 36), cv2.FONT_HERSHEY_SIMPLEX, 0.55, C_TEXT, 1, cv2.LINE_AA) cv2.putText(img, f'CAM {camera_id}', (w - 180, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) def draw_footer(img, w, h, frame_idx, live_tag): bar_h = 28 overlay_rect(img, 0, h - bar_h, w, h, C_PANEL, alpha=0.55) cv2.putText(img, f'{live_tag} | Frame {frame_idx}', (12, h - 9), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) def draw_skeleton_bold(img, kpts): for (a, b), color in zip(SKELETON, SK_COLORS): if a < len(kpts) and b < len(kpts): xa, ya = int(kpts[a][0]), int(kpts[a][1]) xb, yb = int(kpts[b][0]), int(kpts[b][1]) if xa > 0 and ya > 0 and xb > 0 and yb > 0: cv2.line(img, (xa, ya), (xb, yb), color, 3, cv2.LINE_AA) for kp in kpts: x, y = int(kp[0]), int(kp[1]) if x > 0 and y > 0: cv2.circle(img, (x, y), 6, (255, 255, 255), -1, cv2.LINE_AA) cv2.circle(img, (x, y), 6, (40, 40, 40), 2, cv2.LINE_AA) def draw_popups(img, popups, frame_idx): alive = [] for pop in popups: age = frame_idx - pop['born'] if age > POPUP_LIFETIME: continue alive.append(pop) fade = 1.0 - age / POPUP_LIFETIME y = pop['y'] - int(age * 1.8) color = (int(C_GREEN[0] * fade), int(C_GREEN[1] * fade), int(C_GREEN[2] * fade)) cv2.putText(img, pop['text'], (pop['x'], y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2, cv2.LINE_AA) return alive def draw_batch_banner(img, w, batch_num, pulse_remaining): if pulse_remaining <= 0: return text = f'NEW BATCH {batch_num}' font = cv2.FONT_HERSHEY_SIMPLEX (tw, th), _ = cv2.getTextSize(text, font, 0.8, 2) x1, y1 = w // 2 - tw // 2 - 16, 62 x2, y2 = w // 2 + tw // 2 + 16, 62 + th + 20 overlay_rect(img, x1, y1, x2, y2, C_PANEL, alpha=0.7) cv2.rectangle(img, (x1, y1), (x2, y2), C_ACCENT, 2) cv2.putText(img, text, (w // 2 - tw // 2, 62 + th + 4), font, 0.8, C_ACCENT, 2, cv2.LINE_AA) def connect_stream(source, warmup=WARMUP_FRAMES): attempts = 0 while not shutdown_requested: cap = open_capture(source) if not cap.isOpened(): attempts += 1 if MAX_RECONNECT_ATTEMPTS and attempts >= MAX_RECONNECT_ATTEMPTS: raise RuntimeError(f'Cannot open source after {attempts} attempts: {source}') print(f'Cannot open source, retry in {RECONNECT_DELAY_SEC}s...') time.sleep(RECONNECT_DELAY_SEC) continue if warmup > 0 and source.lower().startswith(('rtsp://', 'http://')): warmup_stream(cap, warmup) w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) fps = cap.get(cv2.CAP_PROP_FPS) if not fps or fps <= 1: fps = OUTPUT_FPS return cap, w, h, fps return None, 0, 0, OUTPUT_FPS # ============================================================================= # Main loop # ============================================================================= def run(): global shutdown_requested store = BatchStore( db_path=DB_PATH, state_file=STATE_FILE, camera_name=CAMERA_NAME, object_label=OBJECT_LABEL, cutoff_time=DAILY_CUTOFF_TIME, batch_timeout=BATCH_TIMEOUT_SECONDS, ignore_batch_label_timeout=IGNORE_BATCH_LABEL_TIMEOUT, min_object_per_batch=MIN_OBJECT_PER_BATCH, min_duration_per_batch=MIN_DURATION_PER_BATCH, logger=lambda msg: print(f'[{now_str()}] {msg}'), ) store.start_cutoff_watcher() cross_logger = None if EXPORT_CSV: cross_logger = CsvLogger(CROSS_CSV, ['batch', 'frame', 'timestamp', 'chicken_id']) # Load RKNN model model = RKNNYOLO( model_path=MODEL_PATH, core_mask=CORE_MASK, imgsz=IMGSZ, conf=CONF, num_classes=NUM_CLASSES, score_sigmoid=SCORE_SIGMOID, ) # Class IDs — order comes from RKNN model output (class index) # class index 0 → CLASS_AYAM, index 1 → CLASS_TALENAN (or env-specified) # Use class names in env order: first CLASS_AYAM → id 0, then CLASS_TALENAN → id 1 CLASS_IDS = { os.getenv('CLASS_AYAM', 'ayam'): 0, os.getenv('CLASS_TALENAN', 'talenan'): 1, } ayam_cls = CLASS_IDS[CLASS_AYAM] talenan_cls = CLASS_IDS[CLASS_TALENAN] ayam_tracker = SimpleTracker(max_age=60) talenan_tracker = SimpleTracker(max_age=60) ayam_line_crossed = set() talenan_line_crossed = set() ayam_cross_flash = {} talenan_cross_flash = {} line_pulse = count_pulse = batch_pulse = 0 popups = [] session_start = time.time() frame_idx = 0 video_writer = None cap, w, h, fps = connect_stream(SOURCE) if cap is None: store.shutdown() model.release() return line_x = resolve_line_x(w) print(f'RKNN counter | {w}x{h} @ {fps}fps | line x={line_x} | cross={CROSS_DIRECTION}') print(f'Model: {MODEL_PATH} | imgsz={IMGSZ} | core_mask={CORE_MASK}') print(f'DB: {DB_PATH}') print(f'State: {STATE_FILE}') if RECORD_VIDEO: video_writer = VideoSegmentWriter(OUTPUT_DIR, w, h, fps, VIDEO_SEGMENT_SEC) reconnect_count = 0 while not shutdown_requested: ret, frame = cap.read() if not ret: if not IS_LIVE: break reconnect_count += 1 print(f'Stream dropped (attempt {reconnect_count}), reconnecting in {RECONNECT_DELAY_SEC}s...') cap.release() time.sleep(RECONNECT_DELAY_SEC) cap, w, h, fps = connect_stream(SOURCE) if cap is None: break line_x = resolve_line_x(w) continue now = time.time() elapsed = now - session_start mono = time.monotonic() ayam_crossed_frame = batch_closed_frame = batch_started_frame = False # RKNN inference detections = model(frame) if detections: ayam_dets = [] # list of (cx, xywh_box) talenan_dets = [] ayam_kpts_map = {} # det_idx → keypoints talenan_kpts_map = {} for di, det in enumerate(detections): bbox = det['bbox'] cls_id = det['cls'] cx = (bbox[0] + bbox[2]) / 2.0 x1, y1, x2, y2 = bbox wb, hb = x2 - x1, y2 - y1 box_cxcywh = np.array([cx, (y1 + y2) / 2, wb, hb], dtype=np.float32) if cls_id == talenan_cls: talenan_dets.append((cx, box_cxcywh)) if det['keypoints'] is not None: talenan_kpts_map[len(talenan_dets) - 1] = det['keypoints'] elif cls_id == ayam_cls: ayam_dets.append((cx, box_cxcywh)) if det['keypoints'] is not None: ayam_kpts_map[len(ayam_dets) - 1] = det['keypoints'] # Track ayam — returns (track_id → cx, detection_idx → track_id) ayam_cx_map, ayam_det_to_track = ayam_tracker.update(ayam_dets) # Track talenan talenan_cx_map, talenan_det_to_track = talenan_tracker.update(talenan_dets) # Process talenan crossings for di, (cx, box) in enumerate(talenan_dets): tid = talenan_det_to_track.get(di) if tid is None: continue if tid in talenan_tracker.tracks: if tid in talenan_tracked: prev_cx = talenan_tracked[tid][0] if crossed_line(prev_cx, cx, line_x) and tid not in talenan_line_crossed: talenan_line_crossed.add(tid) if store.record_talenan_crossing(tid): batch_closed_frame = True talenan_cross_flash[tid] = CROSS_FLASH_FRAMES popups.append({ 'x': int(cx) - 20, 'y': int(box[1]), 'born': frame_idx, 'text': 'BATCH CLOSED', }) talenan_tracked[tid] = (cx, mono) # Process ayam crossings for di, (cx, box) in enumerate(ayam_dets): tid = ayam_det_to_track.get(di) if tid is None: continue if tid in ayam_tracker.tracks: if tid in ayam_tracked: prev_cx = ayam_tracked[tid][0] if crossed_line(prev_cx, cx, line_x) and tid not in ayam_line_crossed: ayam_line_crossed.add(tid) _, started_new = store.record_ayam_crossing(tid) if cross_logger: cross_logger.write_row([ store.current_batch_number, frame_idx, datetime.now().isoformat(), tid, ]) ayam_crossed_frame = True if started_new: batch_started_frame = True ayam_cross_flash[tid] = CROSS_FLASH_FRAMES popups.append({ 'x': int(cx) - 12, 'y': int(box[1]), 'born': frame_idx, 'text': '+1', }) ayam_tracked[tid] = (cx, mono) # Draw talenan for di, (cx, box) in enumerate(talenan_dets): tid = talenan_det_to_track.get(di) if tid is None: continue x1 = int(box[0] - box[2] / 2) y1 = int(box[1] - box[3] / 2) x2 = int(box[0] + box[2] / 2) y2 = int(box[1] + box[3] / 2) flash = talenan_cross_flash.get(tid, 0) color = C_GREEN if flash > 0 else C_TALENAN_BOX cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3 if flash > 0 else 2) draw_pill(frame, f'TALENAN {tid}', x1, y1 - 4, color) # Draw ayam for di, (cx, box) in enumerate(ayam_dets): tid = ayam_det_to_track.get(di) if tid is None: continue x1 = int(box[0] - box[2] / 2) y1 = int(box[1] - box[3] / 2) x2 = int(box[0] + box[2] / 2) y2 = int(box[1] + box[3] / 2) flash = ayam_cross_flash.get(tid, 0) color = C_GREEN if flash > 0 else C_AYAM_BOX cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3 if flash > 0 else 2) draw_pill(frame, f'ID {tid}', x1, y1 - 4, color) kpts = ayam_kpts_map.get(di) if kpts is not None: draw_skeleton_bold(frame, kpts) if ayam_crossed_frame: line_pulse = LINE_PULSE_FRAMES count_pulse = COUNT_PULSE_FRAMES if batch_closed_frame: line_pulse = LINE_PULSE_FRAMES if batch_started_frame: batch_pulse = BATCH_PULSE_FRAMES batch_num = store.current_batch_number or 0 batch_count = store.current_batch_count display_total = store.display_total() rate = (display_total / elapsed * 60) if elapsed > 0 else 0.0 draw_elegant_counting_line(frame, line_x, h, line_pulse) draw_hero_count(frame, line_x, h, batch_count, count_pulse) draw_hud(frame, w, batch_num, batch_count, display_total, elapsed, rate, CAMERA_NAME, now_str()) draw_batch_banner(frame, w, batch_num, batch_pulse) draw_footer(frame, w, h, frame_idx, 'LIVE-RKNN' if IS_LIVE else 'FILE-RKNN') popups = draw_popups(frame, popups, frame_idx) for flash_store in (ayam_cross_flash, talenan_cross_flash): for tid in list(flash_store): flash_store[tid] -= 1 if flash_store[tid] <= 0: del flash_store[tid] line_pulse = max(0, line_pulse - 1) count_pulse = max(0, count_pulse - 1) batch_pulse = max(0, batch_pulse - 1) if video_writer is not None: video_writer.write(frame) if LIVE_STREAM_ENABLED and frame_idx % LIVE_STREAM_EVERY_N == 0: try: Path(LIVE_STREAM_FRAME_PATH).parent.mkdir(parents=True, exist_ok=True) _, jpeg = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, LIVE_STREAM_QUALITY]) with open(LIVE_STREAM_FRAME_PATH, 'wb') as f: f.write(jpeg.tobytes()) except Exception: pass frame_idx += 1 if frame_idx % FLUSH_EVERY_N_FRAMES == 0: print( f'[{now_str()}] Frame {frame_idx} | Batch {batch_num}: {batch_count} ' f'| Total: {display_total} | Uptime {elapsed / 3600:.2f}h' ) prune_stale_tracks(ayam_tracked, mono) prune_stale_tracks(talenan_tracked, mono) cap.release() if video_writer is not None: video_writer.release() if cross_logger: cross_logger.close() model.release() store.shutdown() print('\n=== Batch Summary (SQLite) ===') print(f'Database: {DB_PATH}') # Tracked state dicts: track_id → (cx, monotonic_time) ayam_tracked = {} talenan_tracked = {} if __name__ == '__main__': run()