commit 945eebfe89418f7479f5d182dc3cee5ee1d91881 Author: dsutanto Date: Thu Jun 25 18:47:34 2026 +0700 First commit diff --git a/.env b/.env new file mode 100644 index 0000000..2db1568 --- /dev/null +++ b/.env @@ -0,0 +1,43 @@ +OUTPUT_DIR=/opt/jetson-counter +DB_PATH=/opt/jetson-counter/jetson_counter.db +STATE_FILE=/opt/jetson-counter/current_batch.json +SOURCE=rtsp://user:pass@192.168.0.100:554/stream1 +MODEL_PATH=/opt/jetson-counter/yolo9t.rknn +CAMERA_NAME=CC1 +OBJECT_LABEL=ayam-potong +CLASS_AYAM=ayam +CLASS_TALENAN=talenan +LINE_X= +LINE_X_FRAC=0.5 +CROSS_DIRECTION=rtl +IMGSZ=320 +HALF=false +CONF=0.3 +DEVICE=0 +CORE_MASK=1 +NUM_CLASSES=2 +SCORE_SIGMOID=false +DAILY_CUTOFF_TIME=20:00 +BATCH_TIMEOUT_SECONDS=300 +IGNORE_BATCH_LABEL_TIMEOUT_SECONDS=30 +MIN_OBJECT_PER_BATCH=60 +MIN_DURATION_PER_BATCH=60 +EXPORT_CSV=true +CROSS_CSV=/opt/jetson-counter/batch_crossings.csv +WARMUP_FRAMES=30 +RECONNECT_DELAY_SEC=3 +MAX_RECONNECT_ATTEMPTS=0 +FLUSH_EVERY_N_FRAMES=100 +TRACKED_PRUNE_SEC=300 +RECORD_VIDEO=false +VIDEO_SEGMENT_SEC=3600 +OUTPUT_FPS=15 + +LIVE_STREAM_ENABLED=false +LIVE_STREAM_FRAME_PATH=/dev/shm/jetson-counter/live_frame.jpg +LIVE_STREAM_QUALITY=75 +LIVE_STREAM_EVERY_N=2 +OPENCV_FFMPEG_CAPTURE_OPTIONS=rtsp_transport +tcp | fflags +nobuffer | flags +low_delay diff --git a/.env.sample b/.env.sample new file mode 100644 index 0000000..9f45379 --- /dev/null +++ b/.env.sample @@ -0,0 +1,87 @@ +# Copy to .env and adjust values for your deployment + +# Paths +OUTPUT_DIR=/opt/jetson-counter +DB_PATH=/opt/jetson-counter/jetson_counter.db +STATE_FILE=/opt/jetson-counter/current_batch.json + +# Camera stream (rtsp://, http://, or file path for offline testing) +SOURCE=rtsp://user:pass@192.168.0.100:554/stream1 + +# RKNN model +MODEL_PATH=/opt/jetson-counter/yolo9t.rknn + +# Identity +CAMERA_NAME=CC1 +OBJECT_LABEL=ayam-potong +CLASS_AYAM=ayam +CLASS_TALENAN=talenan + +# Counting line: pixel position (empty = auto from LINE_X_FRAC) +LINE_X= +LINE_X_FRAC=0.5 +CROSS_DIRECTION=rtl + +# Model input size +IMGSZ=320 + +# Inference (ignored by RKNN scripts; used by Jetson/TensorRT variant) +HALF=false + +# Confidence threshold for detections +CONF=0.3 + +# TensorRT device index (ignored by RKNN scripts; used by Jetson variant) +DEVICE=0 + +# RKNN NPU core mask: 1=core0, 2=core1, 3=dual, 7=all +CORE_MASK=1 + +# YOLO decoder: number of classes +NUM_CLASSES=2 + +# Set to true if model outputs raw logits instead of sigmoid probabilities +SCORE_SIGMOID=false + +# ByteTrack parameters (counter_live_rknn_bytetrack.py only) +TRACK_HIGH_THRESH=0.5 +TRACK_LOW_THRESH=0.1 +TRACK_MATCH_THRESH=0.8 +TRACK_BUFFER=30 +TRACK_MIN_HITS=3 + +# Batch / cutoff +DAILY_CUTOFF_TIME=20:00 +BATCH_TIMEOUT_SECONDS=300 +IGNORE_BATCH_LABEL_TIMEOUT_SECONDS=30 +MIN_OBJECT_PER_BATCH=60 +MIN_DURATION_PER_BATCH=60 + +# CSV export +EXPORT_CSV=true +CROSS_CSV=/opt/jetson-counter/batch_crossings.csv + +# Stream connection +WARMUP_FRAMES=30 +RECONNECT_DELAY_SEC=3 +MAX_RECONNECT_ATTEMPTS=0 + +# Health logging interval +FLUSH_EVERY_N_FRAMES=100 + +# Stale track pruning (seconds) +TRACKED_PRUNE_SEC=300 + +# Video recording +RECORD_VIDEO=false +VIDEO_SEGMENT_SEC=3600 +OUTPUT_FPS=15 + +# Live stream (writes JPEG snapshot to disk for nginx) +LIVE_STREAM_ENABLED=false +LIVE_STREAM_FRAME_PATH=/dev/shm/jetson-counter/live_frame.jpg +LIVE_STREAM_QUALITY=75 +LIVE_STREAM_EVERY_N=2 + +# OpenCV FFmpeg backend options (semicolon/pipe separated) +OPENCV_FFMPEG_CAPTURE_OPTIONS=rtsp_transport;tcp|fflags;nobuffer|flags;low_delay diff --git a/DEPLOY.md b/DEPLOY.md new file mode 100644 index 0000000..94cc2a4 --- /dev/null +++ b/DEPLOY.md @@ -0,0 +1,122 @@ +# Edge Jetson Deploy + +Production counter: **direct LAN RTSP** + **YOLO11n TensorRT** + SQLite batch store. +Replaces MQTT `frigate-counter` on the edge Jetson. + +## Quick install + +```bash +# 1. Copy this folder to Jetson +sudo mkdir -p /opt/jetson-counter +sudo cp -r jetson-counter/* /opt/jetson-counter/ +sudo chown -R jetson:jetson /opt/jetson-counter + +# 2. Configure +cd /opt/jetson-counter +cp config.env.example .env +nano .env # SOURCE, MODEL_PATH, CAMERA_NAME, etc. +sed -i 's/\r$//' .env + +# 3. Venv + services +chmod +x setup-venv.sh install-services.sh +sudo ./setup-venv.sh +sudo ./install-services.sh +``` + +Dashboard: `http://:5000` + +--- + +## YOLO11n TensorRT engine (one-time) + +On the Jetson (must match `IMGSZ` / `HALF` in `.env`): + +```bash +source /opt/jetson-counter/venv/bin/activate +export PYTHONNOUSERSITE=1 + +yolo export model=/media/jetson/DATA/yolo11n.pt format=engine half=True imgsz=416 device=0 +``` + +Verify classes: + +```bash +PYTHONNOUSERSITE=1 python -c " +from ultralytics import YOLO +m = YOLO('/media/jetson/DATA/yolo11n.engine') +print(m.names) +" +``` + +Expect `ayam` and `talenan`. + +--- + +## Direct camera RTSP + +Set in `.env`: + +```env +SOURCE=rtsp://user:pass@192.168.x.x:554/stream1 +``` + +Test before install: + +```bash +ffplay -rtsp_transport tcp -t 5 "$SOURCE" +nc -zv 554 +``` + +--- + +## Cutover from MQTT frigate-counter + +`install-services.sh` automatically: + +1. Disables `frigate-counter` and `frigate-counter-dashboard` +2. Enables `jetson-counter` + `jetson-counter-dashboard` + +Archive old DB (optional): + +```bash +sudo cp /opt/frigate-counter/frigate_counter.db ~/frigate_counter.db.backup +``` + +--- + +## Validation checklist + +```bash +sudo systemctl is-active jetson-counter jetson-counter-dashboard +PYTHONNOUSERSITE=1 /opt/jetson-counter/venv/bin/python -c "import torch; print('cuda', torch.cuda.is_available())" +sudo journalctl -u jetson-counter -n 20 --no-pager +``` + +Good signs: + +- `Stream ready!` +- `Loaded engine size: ... MiB` +- `Frame 100 | Batch ...` + +--- + +## Logs & restart + +```bash +sudo journalctl -u jetson-counter -f +sudo systemctl restart jetson-counter # after .env change +``` + +--- + +## JetPack 6.0 torch wheel + +If `setup-venv.sh` fails on torch URL, list wheels: + +```bash +curl -s https://developer.download.nvidia.com/compute/redist/jp/v60/pytorch/ | grep cp310 +``` + +Set `TORCH_WHEEL_URL=...` when running `setup-venv.sh`. + +See also [jetson-counter-dev/GO_LIVE_TROUBLESHOOT.md](../jetson-counter-dev/GO_LIVE_TROUBLESHOOT.md) for torchvision and RTSP issues. diff --git a/README.md b/README.md new file mode 100644 index 0000000..193fd06 --- /dev/null +++ b/README.md @@ -0,0 +1,59 @@ +# Jetson Edge Counter (Production) + +RTSP + YOLO TensorRT line-crossing counter for edge Jetson. Replaces MQTT `frigate-counter` on site. + +## Architecture + +- **Input:** Direct LAN camera RTSP (low latency) +- **Inference:** YOLO11n `.engine` (TensorRT) on Jetson GPU +- **Logic:** Line crossing (`ayam` count, `talenan` closes batch) via `batch_store.py` +- **Output:** `jetson_counter.db` + `current_batch.json` +- **Dashboard:** Flask on port **5000** + +Batch lifecycle: talenan closes batch → idle until next ayam line cross (count starts at 1). + +## Deploy + +See **[DEPLOY.md](DEPLOY.md)**. + +| Item | Default | +|------|---------| +| Install path | `/opt/jetson-counter` | +| Venv | `/opt/jetson-counter/venv` | +| DB | `/opt/jetson-counter/jetson_counter.db` | +| Dashboard | `http://:5000` | +| Cutoff | `20:00` | + +## Commands + +| Command | Purpose | +|---------|---------| +| `sudo systemctl status jetson-counter` | Counter running? | +| `sudo journalctl -u jetson-counter -f` | Live logs | +| `sudo systemctl restart jetson-counter` | After `.env` change | +| `sudo ./uninstall-services.sh` | Remove services | + +## Key env vars + +| Variable | Purpose | +|----------|---------| +| `SOURCE` | Direct camera RTSP URL | +| `MODEL_PATH` | `.engine` file path | +| `IMGSZ` / `HALF` | Must match engine export | +| `CROSS_DIRECTION` | `rtl` (default), `ltr`, or `both` | +| `LINE_X` / `LINE_X_FRAC` | Counting line position | + +## Files + +| File | Purpose | +|------|---------| +| `counter_live.py` | RTSP + YOLO + line crossing | +| `batch_store.py` | SQLite persistence | +| `counter_dashboard.py` | Flask UI | +| `config.env.example` | Env template | +| `jetson-counter.service` | Counter systemd unit | +| `install-services.sh` | Install + disable legacy MQTT counter | + +## Dev stack + +Lab / comparison: [`jetson-counter-dev/`](../jetson-counter-dev/) (port 8081, separate DB). diff --git a/batch_store.py b/batch_store.py new file mode 100644 index 0000000..94ee985 --- /dev/null +++ b/batch_store.py @@ -0,0 +1,379 @@ +""" +Production batch persistence for edge Jetson counter. +Mirrors frigate-counter SQLite schema + current_batch.json contract. +""" +import json +import sqlite3 +import threading +import time +from datetime import datetime, timedelta +from pathlib import Path + + +class BatchStore: + def __init__( + self, + db_path, + state_file, + camera_name, + object_label='ayam-potong', + cutoff_time='20:00', + batch_timeout=300.0, + ignore_batch_label_timeout=30.0, + min_object_per_batch=60, + min_duration_per_batch=60, + carry_ids=50, + logger=print, + ): + self.db_path = db_path + self.state_file = Path(state_file) + self.camera_name = camera_name + self.object_label = object_label + self.cutoff_time_str = cutoff_time + datetime.strptime(cutoff_time, '%H:%M') + + self.batch_timeout = float(batch_timeout) + self.ignore_batch_label_timeout = float(ignore_batch_label_timeout) + self.min_object_per_batch = int(min_object_per_batch) + self.min_duration_per_batch = int(min_duration_per_batch) + self.carry_ids = int(carry_ids) + self.log = logger + + self.state_lock = threading.Lock() + self.batch_timer = None + self.ignore_batch_label = False + self.ignore_batch_label_timer = None + self.previous_state = None + self.shutdown_event = threading.Event() + + Path(db_path).parent.mkdir(parents=True, exist_ok=True) + self.state_file.parent.mkdir(parents=True, exist_ok=True) + + self.db = sqlite3.connect(db_path, check_same_thread=False) + self._init_db() + self.current_state = self._load_state() + self.previous_state = self.current_state + + def _init_db(self): + cur = self.db.cursor() + cur.execute( + """ + CREATE TABLE IF NOT EXISTS batches ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + counting_date TEXT NOT NULL, + batch_number INTEGER NOT NULL, + camera_name TEXT NOT NULL, + object_label TEXT NOT NULL, + count INTEGER NOT NULL, + start_time TEXT NOT NULL, + end_time TEXT NOT NULL, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + UNIQUE(counting_date, batch_number, camera_name, object_label) + ) + """ + ) + cur.execute( + """ + CREATE TABLE IF NOT EXISTS daily_summaries ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + counting_date TEXT NOT NULL, + camera_name TEXT NOT NULL, + object_label TEXT NOT NULL, + total_count INTEGER NOT NULL DEFAULT 0, + total_batches INTEGER NOT NULL DEFAULT 0, + updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + UNIQUE(counting_date, camera_name, object_label) + ) + """ + ) + self.db.commit() + + def get_counting_date(self, dt=None): + if dt is None: + dt = datetime.now() + cutoff = datetime.strptime(self.cutoff_time_str, '%H:%M').time() + if dt.time() < cutoff: + return dt.date().isoformat() + return (dt.date() + timedelta(days=1)).isoformat() + + def _load_state(self): + if not self.state_file.exists(): + return None + try: + with open(self.state_file, 'r', encoding='utf-8') as f: + state = json.load(f) + current_date = self.get_counting_date() + if state.get('counting_date') != current_date: + self.log( + f"State file belongs to previous counting day ({state.get('counting_date')}). " + 'Finalizing before fresh start.' + ) + self._insert_batch( + state['counting_date'], + state['batch_number'], + state['count'], + state['start_time'], + datetime.now().isoformat(), + ) + self.state_file.unlink(missing_ok=True) + return None + self.log( + f"Resumed batch #{state['batch_number']} from {state['start_time']} " + f"with count={state['count']}" + ) + self._reset_batch_timer() + return state + except Exception as exc: + self.log(f'Failed to load state file: {exc}') + return None + + def save_state(self): + if self.current_state is None: + self.state_file.unlink(missing_ok=True) + return + with open(self.state_file, 'w', encoding='utf-8') as f: + json.dump(self.current_state, f, indent=2, ensure_ascii=False) + + def get_next_batch_number(self, counting_date): + cur = self.db.cursor() + cur.execute( + """ + SELECT COALESCE(MAX(batch_number), 0) + FROM batches + WHERE counting_date = ? AND camera_name = ? AND object_label = ? + """, + (counting_date, self.camera_name, self.object_label), + ) + return cur.fetchone()[0] + 1 + + def start_new_batch(self, counting_date): + batch_number = self.get_next_batch_number(counting_date) + now = datetime.now().isoformat() + counted_ids = [] + if self.previous_state is not None: + try: + counted_ids = self.previous_state['counted_event_ids'][-self.carry_ids:] + except (KeyError, TypeError): + counted_ids = [] + self.current_state = { + 'counting_date': counting_date, + 'batch_number': batch_number, + 'count': 0, + 'start_time': now, + 'last_detection_time': now, + 'counted_event_ids': counted_ids, + } + self.save_state() + self.log(f'Started batch #{batch_number} for {counting_date} ({self.object_label})') + + def _reset_batch_timer(self): + if self.batch_timer: + self.batch_timer.cancel() + self.batch_timer = threading.Timer(self.batch_timeout, self._on_batch_timeout) + self.batch_timer.daemon = True + self.batch_timer.start() + + def _on_batch_timeout(self): + self.log(f'Batch inactivity timeout ({self.batch_timeout}s) reached') + self.end_batch(closed_by='timeout') + + def _ignore_batch_label(self): + if not self.ignore_batch_label_timer: + self.ignore_batch_label = True + self.ignore_batch_label_timer = threading.Timer( + self.ignore_batch_label_timeout, self._on_ignore_batch_label_timeout + ) + self.ignore_batch_label_timer.daemon = True + self.ignore_batch_label_timer.start() + self.log( + f'Ignore batch label for {self.ignore_batch_label_timeout}s' + ) + + def _on_ignore_batch_label_timeout(self): + self.ignore_batch_label_timer = None + self.ignore_batch_label = False + self.log('Ignore batch label cooldown finished') + + def record_ayam_crossing(self, track_id): + """Line-cross equivalent of production ayam-potong MQTT event.""" + with self.state_lock: + counting_date = self.get_counting_date() + started_new = False + if self.current_state is None: + self.start_new_batch(counting_date) + started_new = True + elif self.current_state['counting_date'] != counting_date: + self._end_batch_locked(closed_by='cutoff') + self.start_new_batch(counting_date) + started_new = True + + event_key = str(track_id) + if event_key not in self.current_state['counted_event_ids']: + self.current_state['count'] += 1 + self.current_state['counted_event_ids'].append(event_key) + self.log( + f'Counted ayam (track {track_id}) | batch #{self.current_state["batch_number"]} ' + f'total: {self.current_state["count"]}' + ) + + self.current_state['last_detection_time'] = datetime.now().isoformat() + self.save_state() + self._reset_batch_timer() + return self.current_state['count'], started_new + + def record_talenan_crossing(self, track_id): + """Line-cross equivalent of production telenan MQTT batch close.""" + if self.ignore_batch_label: + return False + with self.state_lock: + self._ignore_batch_label() + self._end_batch_locked(closed_by='talenan') + self.log(f'Batch closed by talenan (track {track_id})') + if self.batch_timer: + self.batch_timer.cancel() + self.batch_timer = None + return True + + def end_batch(self, closed_by='manual'): + with self.state_lock: + self._end_batch_locked(closed_by=closed_by) + + def _end_batch_locked(self, closed_by='manual'): + if self.current_state is None: + return False + + self.previous_state = self.current_state + state = self.current_state + + start_time_obj = datetime.fromisoformat(state['start_time']) + end_time_obj = datetime.now() + duration_seconds = (end_time_obj - start_time_obj).total_seconds() + + if (state['count'] < self.min_object_per_batch + or duration_seconds < self.min_duration_per_batch): + self.current_state = None + self.save_state() + if self.batch_timer: + self.batch_timer.cancel() + self.batch_timer = None + self.log( + f'Batch #{state["batch_number"]} discarded ' + f'(count={state["count"]}, duration={duration_seconds:.0f}s)' + ) + return False + + end_time = end_time_obj.isoformat() + try: + self._insert_batch( + state['counting_date'], + state['batch_number'], + state['count'], + state['start_time'], + end_time, + ) + cps = state['count'] / duration_seconds if duration_seconds > 0 else 0 + self.log( + f'Batch #{state["batch_number"]} ended | count={state["count"]} | ' + f'duration={duration_seconds:.0f}s | cps={cps:.3f} | closed_by={closed_by}' + ) + except Exception as exc: + self.log(f'Failed to persist batch: {exc}') + return False + + self.current_state = None + self.save_state() + if self.batch_timer: + self.batch_timer.cancel() + self.batch_timer = None + return True + + def _insert_batch(self, counting_date, batch_number, count, start_time, end_time): + cur = self.db.cursor() + cur.execute( + """ + INSERT INTO batches + (counting_date, batch_number, camera_name, object_label, count, start_time, end_time) + VALUES (?, ?, ?, ?, ?, ?, ?) + """, + (counting_date, batch_number, self.camera_name, self.object_label, count, start_time, end_time), + ) + cur.execute( + """ + INSERT INTO daily_summaries + (counting_date, camera_name, object_label, total_count, total_batches) + VALUES (?, ?, ?, ?, 1) + ON CONFLICT(counting_date, camera_name, object_label) + DO UPDATE SET + total_count = total_count + excluded.total_count, + total_batches = total_batches + excluded.total_batches, + updated_at = CURRENT_TIMESTAMP + """, + (counting_date, self.camera_name, self.object_label, count), + ) + self.db.commit() + + cur.execute( + """ + SELECT total_count, total_batches + FROM daily_summaries + WHERE counting_date = ? AND camera_name = ? AND object_label = ? + """, + (counting_date, self.camera_name, self.object_label), + ) + row = cur.fetchone() + if row: + self.log( + f'Daily totals for {counting_date}: {row[0]} objects across {row[1]} batch(es)' + ) + + def cutoff_watcher_loop(self): + while not self.shutdown_event.is_set(): + time.sleep(60) + with self.state_lock: + if self.current_state is None: + continue + if self.current_state['counting_date'] != self.get_counting_date(): + self.log('Daily cutoff reached – finalizing batch') + self._end_batch_locked(closed_by='cutoff') + + def start_cutoff_watcher(self): + t = threading.Thread(target=self.cutoff_watcher_loop, daemon=True) + t.start() + return t + + @property + def current_batch_number(self): + if self.current_state is None: + return 0 + return self.current_state['batch_number'] + + @property + def current_batch_count(self): + if self.current_state is None: + return 0 + return self.current_state['count'] + + def get_closed_total_for_day(self, counting_date=None): + if counting_date is None: + counting_date = self.get_counting_date() + cur = self.db.cursor() + cur.execute( + """ + SELECT COALESCE(total_count, 0) + FROM daily_summaries + WHERE counting_date = ? AND camera_name = ? AND object_label = ? + """, + (counting_date, self.camera_name, self.object_label), + ) + row = cur.fetchone() + return row[0] if row else 0 + + def display_total(self): + return self.get_closed_total_for_day() + self.current_batch_count + + def shutdown(self): + self.shutdown_event.set() + self.end_batch(closed_by='shutdown') + if self.batch_timer: + self.batch_timer.cancel() + self.db.close() diff --git a/config.env.example b/config.env.example new file mode 100644 index 0000000..ca6fcfd --- /dev/null +++ b/config.env.example @@ -0,0 +1,54 @@ +# Edge RK3588 production counter — copy to .env on device +# cp config.env.example .env && nano .env + +OUTPUT_DIR=/opt/jetson-counter +DB_PATH=/opt/jetson-counter/jetson_counter.db +STATE_FILE=/opt/jetson-counter/current_batch.json + +# Direct LAN camera RTSP (low latency) +SOURCE=rtsp://user:pass@192.168.0.100:554/stream1 +OPENCV_FFMPEG_CAPTURE_OPTIONS=rtsp_transport;tcp|fflags;nobuffer|flags;low_delay + +# RKNN model — export'd from YOLO9t with imgsz=320 +MODEL_PATH=/opt/jetson-counter/yolo9t.rknn +IMGSZ=320 +HALF=false +CONF=0.3 + +# RKNN NPU core mask: 1=core0, 2=core1, 3=core0+core1, 7=all three +CORE_MASK=1 + +# YOLO decoder params (must match model export) +NUM_CLASSES=2 +NUM_KEYPOINTS=9 +REG_MAX=16 +STRIDES=8,16,32 + +CAMERA_NAME=CC1 +OBJECT_LABEL=ayam-potong +CLASS_AYAM=ayam +CLASS_TALENAN=talenan + +# Line crossing: rtl (default) | ltr | both +CROSS_DIRECTION=rtl +LINE_X= +LINE_X_FRAC=0.5 + +DAILY_CUTOFF_TIME=20:00 +CUTOFF_TIME=20:00 +BATCH_TIMEOUT_SECONDS=300 +IGNORE_BATCH_LABEL_TIMEOUT_SECONDS=30 +MIN_OBJECT_PER_BATCH=60 +MIN_DURATION_PER_BATCH=60 + +DASHBOARD_HOST=0.0.0.0 +DASHBOARD_PORT=5000 +SECRET_KEY=change-me-in-production + +EXPORT_CSV=true +RECORD_VIDEO=false + +LIVE_STREAM_ENABLED=false +LIVE_STREAM_FRAME_PATH=/dev/shm/jetson-counter/live_frame.jpg +LIVE_STREAM_QUALITY=75 +LIVE_STREAM_EVERY_N=2 diff --git a/counter_dashboard.py b/counter_dashboard.py new file mode 100644 index 0000000..1596453 --- /dev/null +++ b/counter_dashboard.py @@ -0,0 +1,409 @@ +#!/usr/bin/env python3 +""" +Edge Jetson production counter dashboard. +Reads jetson_counter.db + current_batch.json from jetson-counter stack. +Default port 5000 (replaces frigate-counter dashboard role). +""" + +import json +import os +import sqlite3 +import csv +import time +from io import StringIO +from datetime import datetime, timedelta + +from flask import Flask, render_template, jsonify, request, Response +from werkzeug.serving import WSGIRequestHandler + +from dotenv import load_dotenv +load_dotenv() + + +app = Flask(__name__, template_folder="templates") +app.config["SECRET_KEY"] = os.getenv("SECRET_KEY", "change-me-in-production") + +_DEFAULT_DIR = "/opt/jetson-counter" +DB_PATH = os.getenv("DB_PATH", f"{_DEFAULT_DIR}/jetson_counter.db") +CURRENT_BATCH_PATH = os.getenv("STATE_FILE", os.getenv("CURRENT_BATCH_PATH", f"{_DEFAULT_DIR}/current_batch.json")) +CUTOFF_TIME = os.getenv("CUTOFF_TIME", os.getenv("DAILY_CUTOFF_TIME", "20:00")) + +LIVE_STREAM_FRAME_PATH = os.getenv("LIVE_STREAM_FRAME_PATH", "/dev/shm/jetson-counter/live_frame.jpg") + +@app.route("/api/live-video") +def api_live_video(): + def generate(): + while True: + try: + with open(LIVE_STREAM_FRAME_PATH, "rb") as f: + jpeg = f.read() + yield (b"--frame\r\n" + b"Content-Type: image/jpeg\r\n\r\n" + jpeg + b"\r\n") + except FileNotFoundError: + time.sleep(1.0) + continue + except Exception: + time.sleep(0.5) + continue + time.sleep(0.05) + return Response(generate(), mimetype="multipart/x-mixed-replace; boundary=frame") + + +def get_db(): + conn = sqlite3.connect(DB_PATH) + conn.row_factory = sqlite3.Row + return conn + + +def get_counting_date(dt=None, cutoff_str=CUTOFF_TIME): + if dt is None: + dt = datetime.now() + cutoff = datetime.strptime(cutoff_str, "%H:%M").time() + if dt.time() < cutoff: + return dt.date().isoformat() + return (dt.date() + timedelta(days=1)).isoformat() + + +@app.route("/") +def index(): + return render_template("dashboard.html") + + +@app.route("/api/current-batch") +def api_current_batch(): + try: + with open(CURRENT_BATCH_PATH, "r") as f: + data = json.load(f) + return jsonify( + { + "success": True, + "counting_date": data.get("counting_date"), + "batch_number": data.get("batch_number"), + "count": data.get("count", 0), + "start_time": data.get("start_time"), + "last_detection_time": data.get("last_detection_time"), + } + ) + except FileNotFoundError: + return jsonify( + { + "success": False, + "error": "No active batch", + "count": 0, + "batch_number": None, + "counting_date": None, + } + ), 200 + except Exception as e: + return jsonify( + { + "success": False, + "error": str(e), + "count": 0, + "batch_number": None, + "counting_date": None, + } + ), 500 + + +@app.route("/api/previous-batch") +def api_previous_batch(): + conn = get_db() + cur = conn.cursor() + cur.execute( + """ + SELECT counting_date, batch_number, count, start_time, end_time, + ROUND((julianday(end_time) - julianday(start_time)) * 24 * 60, 1) as duration_minutes + FROM batches + ORDER BY end_time DESC + LIMIT 1 + """ + ) + row = cur.fetchone() + conn.close() + if row: + return jsonify( + { + "success": True, + "date": row["counting_date"], + "batch_number": row["batch_number"], + "count": row["count"], + "start_time": row["start_time"], + "end_time": row["end_time"], + "duration_minutes": row["duration_minutes"], + } + ) + return jsonify({"success": False, "error": "No previous batch"}), 200 + + +@app.route("/api/summary") +def api_summary(): + conn = get_db() + cur = conn.cursor() + today = get_counting_date() + cur.execute( + """ + SELECT COALESCE(total_count, 0) as total_count, + COALESCE(total_batches, 0) as total_batches + FROM daily_summaries + WHERE counting_date = ? + """, + (today,), + ) + today_row = cur.fetchone() + yesterday = (datetime.strptime(today, "%Y-%m-%d").date() - timedelta(days=1)).isoformat() + cur.execute( + """ + SELECT COALESCE(total_count, 0) as total_count, + COALESCE(total_batches, 0) as total_batches + FROM daily_summaries + WHERE counting_date = ? + """, + (yesterday,), + ) + yesterday_row = cur.fetchone() + cur.execute( + """ + SELECT COALESCE(SUM(total_count), 0) as grand_total, + COALESCE(SUM(total_batches), 0) as grand_batches, + COUNT(DISTINCT counting_date) as total_days + FROM daily_summaries + """ + ) + all_time = cur.fetchone() + cur.execute("SELECT ROUND(AVG(total_count), 1) as avg_per_day FROM daily_summaries") + avg = cur.fetchone() + cur.execute( + """ + SELECT counting_date, total_count + FROM daily_summaries + ORDER BY total_count DESC + LIMIT 1 + """ + ) + best = cur.fetchone() + conn.close() + return jsonify( + { + "today": { + "date": today, + "total_count": today_row["total_count"] if today_row else 0, + "total_batches": today_row["total_batches"] if today_row else 0, + }, + "yesterday": { + "date": yesterday, + "total_count": yesterday_row["total_count"] if yesterday_row else 0, + "total_batches": yesterday_row["total_batches"] if yesterday_row else 0, + }, + "all_time": { + "grand_total": all_time["grand_total"], + "grand_batches": all_time["grand_batches"], + "total_days": all_time["total_days"], + }, + "average_per_day": avg["avg_per_day"] or 0, + "best_day": { + "date": best["counting_date"] if best else None, + "count": best["total_count"] if best else 0, + }, + } + ) + + +@app.route("/api/daily-data") +def api_daily_data(): + days = request.args.get("days", 30, type=int) + date_from = (datetime.now() - timedelta(days=days)).date().isoformat() + conn = get_db() + cur = conn.cursor() + cur.execute( + """ + SELECT counting_date, total_count, total_batches, + ROUND(CAST(total_count AS FLOAT) / total_batches, 1) as avg_per_batch + FROM daily_summaries + WHERE counting_date >= ? + ORDER BY counting_date ASC + """, + (date_from,), + ) + daily_data = [ + { + "date": row["counting_date"], + "total_count": row["total_count"], + "total_batches": row["total_batches"], + "avg_per_batch": row["avg_per_batch"] or 0, + } + for row in cur.fetchall() + ] + conn.close() + return jsonify(daily_data) + + +@app.route("/api/day-detail/") +def api_day_detail(date): + conn = get_db() + cur = conn.cursor() + cur.execute( + """ + SELECT batch_number, count, start_time, end_time, + ROUND((julianday(end_time) - julianday(start_time)) * 24 * 60, 1) as duration_minutes + FROM batches + WHERE counting_date = ? + ORDER BY batch_number ASC + """, + (date,), + ) + batches = [] + total_duration = 0 + for row in cur.fetchall(): + duration = row["duration_minutes"] or 0 + total_duration += duration + batches.append( + { + "batch_number": row["batch_number"], + "count": row["count"], + "start_time": row["start_time"], + "end_time": row["end_time"], + "duration_minutes": duration, + } + ) + cur.execute( + """ + SELECT total_count, total_batches + FROM daily_summaries + WHERE counting_date = ? + """, + (date,), + ) + summary = cur.fetchone() + conn.close() + return jsonify( + { + "date": date, + "total_count": summary["total_count"] if summary else 0, + "total_batches": summary["total_batches"] if summary else 0, + "total_duration_minutes": round(total_duration, 1), + "avg_duration_minutes": round(total_duration / len(batches), 1) if batches else 0, + "batches": batches, + } + ) + + +@app.route("/api/recent-batches") +def api_recent_batches(): + limit = request.args.get("limit", 10, type=int) + conn = get_db() + cur = conn.cursor() + cur.execute( + """ + SELECT counting_date, batch_number, count, start_time, end_time, + ROUND((julianday(end_time) - julianday(start_time)) * 24 * 60, 1) as duration_minutes + FROM batches + ORDER BY end_time DESC + LIMIT ? + """, + (limit,), + ) + batches = [ + { + "date": row["counting_date"], + "batch_number": row["batch_number"], + "count": row["count"], + "start_time": row["start_time"], + "end_time": row["end_time"], + "duration_minutes": row["duration_minutes"] or 0, + } + for row in cur.fetchall() + ] + conn.close() + return jsonify(batches) + + +@app.route("/api/available-dates") +def api_available_dates(): + conn = get_db() + cur = conn.cursor() + cur.execute( + """ + SELECT counting_date, total_count, total_batches + FROM daily_summaries + ORDER BY counting_date DESC + """ + ) + dates = [ + { + "date": row["counting_date"], + "total_count": row["total_count"], + "total_batches": row["total_batches"], + } + for row in cur.fetchall() + ] + conn.close() + return jsonify(dates) + + +@app.route("/api/export-daily-csv") +def export_daily_csv(): + days = request.args.get("days", 30, type=int) + date_from = (datetime.now() - timedelta(days=days)).date().isoformat() + conn = get_db() + cur = conn.cursor() + cur.execute( + """ + SELECT counting_date, total_count, total_batches, + ROUND(CAST(total_count AS FLOAT) / NULLIF(total_batches, 0), 1) as avg_per_batch + FROM daily_summaries + WHERE counting_date >= ? + ORDER BY counting_date ASC + """, + (date_from,), + ) + output = StringIO() + writer = csv.writer(output) + writer.writerow(["Date", "Total Count", "Total Batches", "Avg per Batch"]) + for row in cur.fetchall(): + writer.writerow([row["counting_date"], row["total_count"], row["total_batches"], row["avg_per_batch"] or 0]) + conn.close() + filename = f"daily_records_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv" + return Response( + output.getvalue(), + mimetype="text/csv", + headers={"Content-Disposition": f"attachment; filename={filename}"}, + ) + + +@app.route("/api/export-day-csv/") +def export_day_csv(date): + conn = get_db() + cur = conn.cursor() + cur.execute( + """ + SELECT batch_number, count, start_time, end_time, + ROUND((julianday(end_time) - julianday(start_time)) * 24 * 60, 1) as duration_minutes + FROM batches + WHERE counting_date = ? + ORDER BY batch_number ASC + """, + (date,), + ) + output = StringIO() + writer = csv.writer(output) + writer.writerow(["Batch Number", "Count", "Start Time", "End Time", "Duration (min)"]) + for row in cur.fetchall(): + writer.writerow([row["batch_number"], row["count"], row["start_time"], row["end_time"], row["duration_minutes"] or 0]) + conn.close() + return Response( + output.getvalue(), + mimetype="text/csv", + headers={"Content-Disposition": f"attachment; filename=day_detail_{date}.csv"}, + ) + + +if __name__ == "__main__": + WSGIRequestHandler.protocol_version = "HTTP/1.1" + port = int(os.getenv("DASHBOARD_PORT", "5000")) + host = os.getenv("DASHBOARD_HOST", "0.0.0.0") + debug = os.getenv("FLASK_DEBUG", "false").lower() == "true" + print(f"Jetson counter dashboard at http://{host}:{port}") + print(f"DB: {DB_PATH}") + print(f"State: {CURRENT_BATCH_PATH}") + app.run(host=host, port=port, debug=debug) diff --git a/counter_live.py b/counter_live.py new file mode 100644 index 0000000..6a60452 --- /dev/null +++ b/counter_live.py @@ -0,0 +1,559 @@ +""" +Edge production live counter — RTSP + YOLO TensorRT + line crossing. +Replaces MQTT frigate-counter on Jetson with local LAN camera inference. +""" +from ultralytics import YOLO +import cv2 +import csv +import numpy as np +import os +import signal +import time +from datetime import datetime +from pathlib import Path + +from dotenv import load_dotenv +load_dotenv() + +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', '/media/jetson/DATA/yolo11n.engine') +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', '416')) +HALF = os.getenv('HALF', 'true').lower() == 'true' +CONF = float(os.getenv('CONF', '0.3')) +DEVICE = int(os.getenv('DEVICE', '0')) +TRACKER = os.getenv('TRACKER', 'bytetrack.yaml') + +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) + + +def resolve_class_ids(names): + name_to_id = {v: k for k, v in names.items()} + missing = [n for n in (CLASS_AYAM, CLASS_TALENAN) if n not in name_to_id] + if missing: + raise ValueError(f'Model missing classes {missing}. Available: {list(names.values())}') + return name_to_id[CLASS_AYAM], name_to_id[CLASS_TALENAN] + + +def box_cx(box): + return (int(box[0]) + int(box[2])) // 2 + + +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 + + +def run(): + global shutdown_requested + + store = BatchStore( + db_path=DB_PATH, + state_file=STATE_FILE, + camera_name=CAMERA_NAME, + object_label=OBJECT_LABEL, + cutoff_time=DAILY_CUTOFF_TIME, + batch_timeout=BATCH_TIMEOUT_SECONDS, + ignore_batch_label_timeout=IGNORE_BATCH_LABEL_TIMEOUT, + min_object_per_batch=MIN_OBJECT_PER_BATCH, + min_duration_per_batch=MIN_DURATION_PER_BATCH, + logger=lambda msg: print(f'[{now_str()}] {msg}'), + ) + store.start_cutoff_watcher() + + cross_logger = None + if EXPORT_CSV: + cross_logger = CsvLogger(CROSS_CSV, ['batch', 'frame', 'timestamp', 'chicken_id']) + + model = YOLO(MODEL_PATH) + ayam_cls, talenan_cls = resolve_class_ids(model.names) + + ayam_tracked = {} + talenan_tracked = {} + 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() + return + + line_x = resolve_line_x(w) + print(f'Jetson counter | {w}x{h} @ {fps}fps | line x={line_x} | cross={CROSS_DIRECTION}') + print(f'Model: {MODEL_PATH} | imgsz={IMGSZ} half={HALF}') + 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 + + results = model.track( + frame, + device=DEVICE, + persist=True, + conf=CONF, + imgsz=IMGSZ, + half=HALF, + tracker=TRACKER, + verbose=False, + ) + r = results[0] + + if r.boxes.id is not None: + ids = r.boxes.id.int().tolist() + boxes = r.boxes.xyxy.tolist() + clss = r.boxes.cls.int().tolist() + kpts_all = r.keypoints.xy.tolist() if r.keypoints else [] + + talenan_items, ayam_items = [], [] + for i, (track_id, box, cls_id) in enumerate(zip(ids, boxes, clss)): + cx = box_cx(box) + x1, y1, x2, y2 = [int(v) for v in box] + kpts = kpts_all[i] if i < len(kpts_all) else None + item = (track_id, cx, x1, y1, x2, y2, kpts) + if cls_id == talenan_cls: + talenan_items.append(item) + elif cls_id == ayam_cls: + ayam_items.append(item) + + for track_id, cx, x1, y1, x2, y2, _ in talenan_items: + if track_id in talenan_tracked: + prev_cx, _ = talenan_tracked[track_id] + if crossed_line(prev_cx, cx, line_x) and track_id not in talenan_line_crossed: + talenan_line_crossed.add(track_id) + if store.record_talenan_crossing(track_id): + batch_closed_frame = True + talenan_cross_flash[track_id] = CROSS_FLASH_FRAMES + popups.append({'x': cx - 20, 'y': (y1 + y2) // 2, 'born': frame_idx, 'text': 'BATCH CLOSED'}) + talenan_tracked[track_id] = (cx, mono) + + for track_id, cx, x1, y1, x2, y2, kpts in ayam_items: + if track_id in ayam_tracked: + prev_cx, _ = ayam_tracked[track_id] + if crossed_line(prev_cx, cx, line_x) and track_id not in ayam_line_crossed: + ayam_line_crossed.add(track_id) + _, started_new = store.record_ayam_crossing(track_id) + if cross_logger: + cross_logger.write_row([ + store.current_batch_number, frame_idx, + datetime.now().isoformat(), track_id, + ]) + ayam_crossed_frame = True + if started_new: + batch_started_frame = True + ayam_cross_flash[track_id] = CROSS_FLASH_FRAMES + popups.append({'x': cx - 12, 'y': (y1 + y2) // 2, 'born': frame_idx, 'text': '+1'}) + ayam_tracked[track_id] = (cx, mono) + + for track_id, cx, x1, y1, x2, y2, _ in talenan_items: + flash = talenan_cross_flash.get(track_id, 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 {track_id}', x1, y1 - 4, color) + + for track_id, cx, x1, y1, x2, y2, kpts in ayam_items: + flash = ayam_cross_flash.get(track_id, 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 {track_id}', x1, y1 - 4, color) + 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' if IS_LIVE else 'FILE') + 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() + store.shutdown() + + print('\n=== Batch Summary (SQLite) ===') + print(f'Database: {DB_PATH}') + + +if __name__ == '__main__': + run() diff --git a/counter_live_rknn.py b/counter_live_rknn.py new file mode 100644 index 0000000..4d98bd9 --- /dev/null +++ b/counter_live_rknn.py @@ -0,0 +1,868 @@ +""" +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() diff --git a/counter_live_rknn_bytetrack.py b/counter_live_rknn_bytetrack.py new file mode 100644 index 0000000..f905e17 --- /dev/null +++ b/counter_live_rknn_bytetrack.py @@ -0,0 +1,1086 @@ +""" +Edge production live counter — RTSP + YOLO RKNN + ByteTrack + line crossing. +Runs on RK3588 hardware with RKNN model (320×320 input). +Uses ByteTrack (Kalman filter + two-stage IoU association) for tracking. +""" +import numpy as np +import cv2 +import csv +import os +import signal +import time +from datetime import datetime +from pathlib import Path + +from dotenv import load_dotenv +load_dotenv() + +from rknnlite.api import RKNNLite +from batch_store import BatchStore + +# --- config (override via env / .env) --- +OUTPUT_DIR = os.getenv('OUTPUT_DIR', '/opt/jetson-counter') +DB_PATH = os.getenv('DB_PATH', f'{OUTPUT_DIR}/jetson_counter.db') +STATE_FILE = os.getenv('STATE_FILE', f'{OUTPUT_DIR}/current_batch.json') +SOURCE = os.getenv('SOURCE', 'rtsp://user:pass@192.168.0.100:554/stream1') +MODEL_PATH = os.getenv('MODEL_PATH', '/opt/jetson-counter/yolo11n.rknn') +CAMERA_NAME = os.getenv('CAMERA_NAME', 'CC1') +OBJECT_LABEL = os.getenv('OBJECT_LABEL', 'ayam-potong') +CLASS_AYAM = os.getenv('CLASS_AYAM', 'ayam') +CLASS_TALENAN = os.getenv('CLASS_TALENAN', 'talenan') + +LINE_X = int(os.getenv('LINE_X')) if os.getenv('LINE_X') else None +LINE_X_FRAC = float(os.getenv('LINE_X_FRAC', '0.5')) +CROSS_DIRECTION = os.getenv('CROSS_DIRECTION', 'rtl').lower() + +IMGSZ = int(os.getenv('IMGSZ', '320')) +HALF = os.getenv('HALF', 'false').lower() == 'true' +CONF = float(os.getenv('CONF', '0.3')) +DEVICE = int(os.getenv('DEVICE', '0')) + +# RKNN NPU core mask +CORE_MASK = int(os.getenv('CORE_MASK', '1')) + +# YOLO decoder config +NUM_CLASSES = int(os.getenv('NUM_CLASSES', '2')) +SCORE_SIGMOID = os.getenv('SCORE_SIGMOID', 'false').lower() == 'true' + +# ByteTrack settings +TRACK_HIGH_THRESH = float(os.getenv('TRACK_HIGH_THRESH', '0.5')) +TRACK_LOW_THRESH = float(os.getenv('TRACK_LOW_THRESH', '0.1')) +TRACK_MATCH_THRESH = float(os.getenv('TRACK_MATCH_THRESH', '0.8')) +TRACK_BUFFER = int(os.getenv('TRACK_BUFFER', '30')) +TRACK_MIN_HITS = int(os.getenv('TRACK_MIN_HITS', '3')) + +DAILY_CUTOFF_TIME = os.getenv('DAILY_CUTOFF_TIME', '20:00') +BATCH_TIMEOUT_SECONDS = float(os.getenv('BATCH_TIMEOUT_SECONDS', '300')) +IGNORE_BATCH_LABEL_TIMEOUT = float(os.getenv('IGNORE_BATCH_LABEL_TIMEOUT_SECONDS', '30')) +MIN_OBJECT_PER_BATCH = int(os.getenv('MIN_OBJECT_PER_BATCH', '60')) +MIN_DURATION_PER_BATCH = int(os.getenv('MIN_DURATION_PER_BATCH', '60')) + +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) + + +# ============================================================================= +# IoU helpers (xyxy format) +# ============================================================================= + +def _ious_xyxy(boxes_a, boxes_b): + """Pairwise IoU: (N,4) vs (M,4) → (N,M) matrix.""" + n, m = len(boxes_a), len(boxes_b) + if n == 0 or m == 0: + return np.zeros((n, m), dtype=np.float32) + xx1 = np.maximum(boxes_a[:, None, 0], boxes_b[None, :, 0]) + yy1 = np.maximum(boxes_a[:, None, 1], boxes_b[None, :, 1]) + xx2 = np.minimum(boxes_a[:, None, 2], boxes_b[None, :, 2]) + yy2 = np.minimum(boxes_a[:, None, 3], boxes_b[None, :, 3]) + iw = np.maximum(0.0, xx2 - xx1) + ih = np.maximum(0.0, yy2 - yy1) + inter = iw * ih + area_a = (boxes_a[:, 2] - boxes_a[:, 0]) * (boxes_a[:, 3] - boxes_a[:, 1]) + area_b = (boxes_b[:, 2] - boxes_b[:, 0]) * (boxes_b[:, 3] - boxes_b[:, 1]) + return inter / (area_a[:, None] + area_b[None, :] - inter + 1e-16) + + +def _greedy_match(cost_matrix, threshold=0.3): + """Greedy linear assignment. Returns pairs (row_idx, col_idx).""" + if cost_matrix.size == 0: + return [] + n, m = cost_matrix.shape + flat = [(cost_matrix[i, j], i, j) for i in range(n) for j in range(m)] + flat.sort() + row_used = set() + col_used = set() + pairs = [] + for cost, i, j in flat: + if cost >= threshold: + break + if i in row_used or j in col_used: + continue + row_used.add(i) + col_used.add(j) + pairs.append((i, j)) + return pairs + + +# ============================================================================= +# Kalman filter box tracker (state: x, y, w, h, vx, vy, vw, vh) +# ============================================================================= + +class KalmanBoxTracker: + count = 0 + + def __init__(self, bbox_xyxy): + KalmanBoxTracker.count += 1 + self.track_id = KalmanBoxTracker.count + + x1, y1, x2, y2 = bbox_xyxy + w, h = x2 - x1, y2 - y1 + x, y = x1 + w / 2, y1 + h / 2 + + self.kf = _KalmanFilter() + self.kf.x[:4, 0] = np.array([x, y, w, h], dtype=np.float32) + + self.time_since_update = 0 + self.hits = 1 + self.hit_streak = 1 + self.age = 1 + + def predict(self): + if self.kf.x[6] + self.kf.x[2] <= 0: + self.kf.x[6] *= 0.0 + self.kf.predict() + self.age += 1 + self.time_since_update += 1 + + def update(self, bbox_xyxy): + self.time_since_update = 0 + self.hits += 1 + self.hit_streak += 1 + + x1, y1, x2, y2 = bbox_xyxy + w, h = x2 - x1, y2 - y1 + x, y = x1 + w / 2, y1 + h / 2 + self.kf.update(np.array([x, y, w, h], dtype=np.float32)) + + def get_state(self): + """Returns xyxy bbox from Kalman state.""" + xx = self.kf.x[:4, 0] + x, y, w, h = xx[0], xx[1], xx[2], xx[3] + x1 = x - w / 2 + y1 = y - h / 2 + x2 = x + w / 2 + y2 = y + h / 2 + return np.array([x1, y1, x2, y2], dtype=np.float32) + + def get_cx(self): + return float(self.kf.x[0, 0]) + + +class _KalmanFilter: + """8-state constant-velocity Kalman filter for bounding box tracking.""" + def __init__(self): + ndim, dt = 4, 1.0 + self.motion_mat = np.eye(2 * ndim, 2 * ndim, dtype=np.float32) + for i in range(ndim): + self.motion_mat[i, ndim + i] = dt + self.update_mat = np.eye(ndim, 2 * ndim, dtype=np.float32) + + self._std_weight_position = 1.0 / 20 + self._std_weight_velocity = 1.0 / 160 + + self.x = np.zeros((8, 1), dtype=np.float32) + self.P = np.eye(8, dtype=np.float32) * 10.0 + + def predict(self): + std_pos = [ + self._std_weight_position * self.x[2], + self._std_weight_position * self.x[3], + self._std_weight_position * self.x[2], + self._std_weight_position * self.x[3], + ] + std_vel = [ + self._std_weight_velocity * self.x[2], + self._std_weight_velocity * self.x[3], + self._std_weight_velocity * self.x[2], + self._std_weight_velocity * self.x[3], + ] + Q = np.diag(np.square(np.concatenate([std_pos, std_vel]))) + + self.x = self.motion_mat @ self.x + self.P = self.motion_mat @ self.P @ self.motion_mat.T + Q + + def update(self, z): + R = np.diag(np.square([ + self._std_weight_position * z[2], + self._std_weight_position * z[3], + self._std_weight_position * z[2], + self._std_weight_position * z[3], + ])) + + H = self.update_mat + S = H @ self.P @ H.T + R + K = self.P @ H.T @ np.linalg.inv(S) + + y = z.reshape(4, 1) - H @ self.x + self.x = self.x + K @ y + I_KH = np.eye(8) - K @ H + self.P = I_KH @ self.P @ I_KH.T + K @ R @ K.T + + +# ============================================================================= +# ByteTrack multi-object tracker +# ============================================================================= + +class ByteTracker: + """ByteTrack: two-stage association with Kalman filter prediction.""" + + def __init__(self, track_high_thresh=0.5, track_low_thresh=0.1, + match_thresh=0.8, track_buffer=30, min_hits=3): + self.high_thresh = track_high_thresh + self.low_thresh = track_low_thresh + self.match_thresh = match_thresh + self.track_buffer = track_buffer + self.min_hits = min_hits + self.tracked_tracks = [] + self.lost_tracks = [] + self.removed_tracks = [] + self.frame_id = 0 + + def update(self, boxes_xyxy, scores): + self.frame_id += 1 + + # --- separate detections by score --- + if len(boxes_xyxy) > 0: + remain = scores > self.low_thresh + dets = boxes_xyxy[remain] + det_scores = scores[remain] + is_high = det_scores > self.high_thresh + is_low = ~is_high + else: + 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]) + 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[det_global] = track_pool[ti].track_id + tracked_map[track_pool[ti].track_id] = track_pool[ti].get_cx() + match_pairs_high.append((det_global, ti)) + + unmatched_tracks = [t for t in range(num_tracks) if t not in matched_track_idx] + + # --- second association: low-score ↔ unmatched tracks --- + low_idx = np.where(is_low)[0] + low_dets = dets[is_low] + + if len(low_dets) > 0 and len(unmatched_tracks) > 0: + unmatched_boxes = track_boxes[unmatched_tracks] + iou_mat = _ious_xyxy(low_dets, unmatched_boxes) + cost_mat = 1.0 - iou_mat + matches2 = _greedy_match(cost_mat, threshold=0.5) + + for di, uti in matches2: + det_global = int(low_idx[di]) + pool_idx = unmatched_tracks[uti] + track_pool[pool_idx].update(dets[det_global]) + track_pool[pool_idx].hit_streak = max(1, track_pool[pool_idx].hit_streak) + matched_track_idx.add(pool_idx) + det_to_track[det_global] = track_pool[pool_idx].track_id + tracked_map[track_pool[pool_idx].track_id] = track_pool[pool_idx].get_cx() + + # --- reset hit_streak for unmatched tracks --- + for ti, trk in enumerate(track_pool): + if ti not in matched_track_idx: + trk.hit_streak = 0 + + # --- lifecycle management --- + new_tracked = [] + new_lost = [] + for trk in track_pool: + if trk.time_since_update > self.track_buffer: + self.removed_tracks.append(trk) + elif trk.time_since_update > 0: + new_lost.append(trk) + else: + new_tracked.append(trk) + self.tracked_tracks = new_tracked + self.lost_tracks = new_lost + + # --- confirmed tracks (both tracked and lost) --- + for trk in self.tracked_tracks + self.lost_tracks: + if trk.hit_streak >= self.min_hits or trk.hits >= self.min_hits: + tracked_map.setdefault(trk.track_id, trk.get_cx()) + + for trk in self.lost_tracks: + if trk.hit_streak >= self.min_hits or trk.hits >= self.min_hits: + lost_map[trk.track_id] = trk.get_cx() + + # --- new tracks from unmatched high-score dets --- + high_all = np.where(is_high)[0] + matched_det_ids = set(det_to_track.keys()) + for dg in high_all: + if int(dg) not in matched_det_ids: + trk = KalmanBoxTracker(dets[dg]) + self.tracked_tracks.append(trk) + det_to_track[int(dg)] = trk.track_id + tracked_map[trk.track_id] = trk.get_cx() + + return tracked_map, det_to_track, lost_map + + +# ============================================================================= +# RKNN YOLO wrapper (detect output format: (1, 4+num_classes, N)) +# ============================================================================= + +class RKNNYOLO: + def __init__(self, model_path, core_mask=1, imgsz=320, conf=0.3, iou=0.45, + num_classes=2, num_keypoints=0, score_sigmoid=False): + self.imgsz = imgsz + self.conf = conf + self.iou = iou + self.num_classes = num_classes + self.num_keypoints = num_keypoints + self.score_sigmoid = score_sigmoid + + self.rknn = RKNNLite(verbose=False) + ret = self.rknn.load_rknn(model_path) + if ret != 0: + raise RuntimeError(f'Failed to load RKNN model: {model_path}') + ret = self.rknn.init_runtime(core_mask=core_mask) + if ret != 0: + raise RuntimeError(f'Failed to init RKNN runtime (core_mask={core_mask})') + + try: + sdk_ver = self.rknn.get_sdk_version() + print(f'RKNN SDK version: {sdk_ver}') + except Exception: + pass + + print(f'RKNN model loaded: {model_path} imgsz={imgsz} core_mask={core_mask}') + + def _preprocess(self, frame): + h0, w0 = frame.shape[:2] + scale = min(self.imgsz / h0, self.imgsz / w0) + nh, nw = int(h0 * scale), int(w0 * scale) + resized = cv2.resize(frame, (nw, nh), interpolation=cv2.INTER_LINEAR) + + letterbox = np.full((self.imgsz, self.imgsz, 3), 114, dtype=np.uint8) + dy = (self.imgsz - nh) // 2 + dx = (self.imgsz - nw) // 2 + letterbox[dy:dy + nh, dx:dx + nw] = resized + + rgb = cv2.cvtColor(letterbox, cv2.COLOR_BGR2RGB) + gains = np.array([scale, scale, dy, dx], dtype=np.float32) + return rgb, gains + + def __call__(self, frame): + h0, w0 = frame.shape[:2] + rgb, gains = self._preprocess(frame) + scale, _, pad_y, pad_x = gains + + inp = np.expand_dims(rgb, axis=0) + inp = np.ascontiguousarray(inp.astype(np.uint8)) + + outputs = self.rknn.inference(inputs=[inp]) + + if len(outputs) == 0: + return [] + + out = outputs[0] + out = np.squeeze(out, axis=0) + + if out.shape[0] == self.num_classes + 4: + out = out.T + + boxes_cxcywh = out[:, :4].copy() + cls_raw = out[:, 4:].copy() + + if self.score_sigmoid: + cls_scores = 1.0 / (1.0 + np.exp(-np.clip(cls_raw, -10, 10))) + else: + cls_scores = cls_raw + + boxes_xyxy = np.stack([ + boxes_cxcywh[:, 0] - boxes_cxcywh[:, 2] / 2, + boxes_cxcywh[:, 1] - boxes_cxcywh[:, 3] / 2, + boxes_cxcywh[:, 0] + boxes_cxcywh[:, 2] / 2, + boxes_cxcywh[:, 1] + boxes_cxcywh[:, 3] / 2, + ], axis=1) + + max_scores = cls_scores.max(axis=1) + class_ids = cls_scores.argmax(axis=1) + mask = max_scores > self.conf + + if mask.sum() == 0: + return [] + + bboxes = boxes_xyxy[mask].astype(np.float32) + scores = max_scores[mask].astype(np.float32) + clses = class_ids[mask] + + bboxes[:, 0] = (bboxes[:, 0] - pad_x) / scale + bboxes[:, 1] = (bboxes[:, 1] - pad_y) / scale + bboxes[:, 2] = (bboxes[:, 2] - pad_x) / scale + bboxes[:, 3] = (bboxes[:, 3] - pad_y) / scale + + bboxes[:, 0] = np.clip(bboxes[:, 0], 0, w0) + bboxes[:, 1] = np.clip(bboxes[:, 1], 0, h0) + bboxes[:, 2] = np.clip(bboxes[:, 2], 0, w0) + bboxes[:, 3] = np.clip(bboxes[:, 3], 0, h0) + + detections = [] + for cls_id in range(self.num_classes): + idx = np.where(clses == cls_id)[0] + if len(idx) == 0: + continue + keep = _nms(bboxes[idx], scores[idx], iou_thr=self.iou) + for k in keep: + j = idx[k] + detections.append({ + 'bbox': bboxes[j].tolist(), + 'score': float(scores[j]), + 'cls': int(clses[j]), + 'keypoints': None, + }) + return detections + + def release(self): + self.rknn.release() + + +# ============================================================================= +# Drawing helpers +# ============================================================================= + +def resolve_line_x(frame_width): + if LINE_X is not None: + return LINE_X + if LINE_X_FRAC != 0.5: + return int(frame_width * LINE_X_FRAC) + return frame_width // 2 + + +def crossed_line(prev_cx, cx, line_x, direction=CROSS_DIRECTION): + if direction == 'ltr': + return prev_cx < line_x <= cx + if direction == 'both': + return (prev_cx > line_x >= cx) or (prev_cx < line_x <= cx) + return prev_cx > line_x >= cx + + +def now_str(): + return datetime.now().strftime('%Y-%m-%d %H:%M:%S') + + +def open_capture(source): + if source.lower().startswith(('rtsp://', 'http://')): + os.environ['OPENCV_FFMPEG_CAPTURE_OPTIONS'] = RTSP_FFMPEG_OPTIONS + cap = cv2.VideoCapture(source, cv2.CAP_FFMPEG) + cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) + return cap + + +def warmup_stream(cap, n=WARMUP_FRAMES): + print('Warming up stream...') + for _ in range(n): + cap.read() + print('Stream ready!') + + +def open_video_writer(path, w, h, fps): + return cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*'avc1'), fps, (w, h)) + + +class CsvLogger: + def __init__(self, path, header): + Path(path).parent.mkdir(parents=True, exist_ok=True) + new_file = not Path(path).exists() or Path(path).stat().st_size == 0 + self.file = open(path, 'a', newline='', buffering=1) + self.writer = csv.writer(self.file) + if new_file: + self.writer.writerow(header) + self.file.flush() + + def write_row(self, row): + self.writer.writerow(row) + self.file.flush() + + def close(self): + self.file.close() + + +class VideoSegmentWriter: + def __init__(self, output_dir, w, h, fps, segment_sec): + self.output_dir = Path(output_dir) + self.output_dir.mkdir(parents=True, exist_ok=True) + self.w, self.h, self.fps = w, h, fps + self.segment_sec = segment_sec + self.segment_start = time.monotonic() + self.writer = None + self._open_next() + + def _segment_path(self): + ts = datetime.now().strftime('%Y%m%d_%H%M%S') + return str(self.output_dir / f'live_{ts}.mp4') + + def _open_next(self): + if self.writer is not None: + self.writer.release() + path = self._segment_path() + self.writer = open_video_writer(path, self.w, self.h, self.fps) + self.segment_start = time.monotonic() + print(f'Recording segment: {path}') + + def write(self, frame): + if time.monotonic() - self.segment_start >= self.segment_sec: + self._open_next() + self.writer.write(frame) + + def release(self): + if self.writer is not None: + self.writer.release() + + +def prune_stale_tracks(tracked, now_mono): + stale = [tid for tid, (_, ts) in tracked.items() if now_mono - ts > TRACKED_PRUNE_SEC] + for tid in stale: + del tracked[tid] + + +def overlay_rect(img, x1, y1, x2, y2, color, alpha=0.65): + x1, y1 = max(0, x1), max(0, y1) + x2, y2 = min(img.shape[1], x2), min(img.shape[0], y2) + if x2 <= x1 or y2 <= y1: + return + roi = img[y1:y2, x1:x2] + patch = np.full_like(roi, color, dtype=np.uint8) + cv2.addWeighted(patch, alpha, roi, 1 - alpha, 0, roi) + + +def draw_pill(img, text, x, y, bg, fg=C_TEXT, font_scale=0.45, pad_x=6, pad_y=4): + font = cv2.FONT_HERSHEY_SIMPLEX + (tw, th), baseline = cv2.getTextSize(text, font, font_scale, 1) + x1, y1 = x, y - th - pad_y + x2, y2 = x + tw + pad_x * 2, y + baseline + pad_y + cv2.rectangle(img, (x1, y1), (x2, y2), bg, -1) + cv2.rectangle(img, (x1, y1), (x2, y2), C_BORDER, 1) + cv2.putText(img, text, (x + pad_x, y), font, font_scale, fg, 1, cv2.LINE_AA) + + +def draw_elegant_counting_line(img, line_x, h, pulse_remaining=0): + strength = pulse_remaining / max(LINE_PULSE_FRAMES, 1) + glow_alpha = 0.12 + 0.18 * strength + for offset in (14, 9, 5): + color = tuple(int(c * glow_alpha) for c in C_LINE_GLOW) + cv2.line(img, (line_x - offset, 0), (line_x - offset, h), color, 1, cv2.LINE_AA) + cv2.line(img, (line_x + offset, 0), (line_x + offset, h), color, 1, cv2.LINE_AA) + dash_len, gap = 18, 12 + y = 0 + while y < h: + y_end = min(y + dash_len, h) + cv2.line(img, (line_x, y), (line_x, y_end), C_LINE_CORE, 2, cv2.LINE_AA) + y += dash_len + gap + cv2.putText(img, 'COUNT LINE', (line_x - 46, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.42, C_LINE_CORE, 1, cv2.LINE_AA) + + +def draw_hero_count(img, line_x, h, count, pulse_remaining=0): + text = str(count) + font = cv2.FONT_HERSHEY_SIMPLEX + boost = 0.35 * (pulse_remaining / max(COUNT_PULSE_FRAMES, 1)) + font_scale, thickness = 1.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 '\u2014' + cv2.putText(img, batch_label, (16, 44), cv2.FONT_HERSHEY_SIMPLEX, 0.9, C_ACCENT, 2, cv2.LINE_AA) + cv2.putText(img, 'COUNT', (100, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) + cv2.putText(img, str(batch_count), (100, 44), cv2.FONT_HERSHEY_SIMPLEX, 0.9, C_GREEN, 2, cv2.LINE_AA) + cv2.putText(img, 'TOTAL', (190, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) + cv2.putText(img, str(total_ayam), (190, 44), cv2.FONT_HERSHEY_SIMPLEX, 0.7, C_TEXT, 1, cv2.LINE_AA) + cv2.putText(img, 'UPTIME', (280, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) + cv2.putText(img, f'{elapsed_sec / 3600:.1f}h', (280, 44), cv2.FONT_HERSHEY_SIMPLEX, 0.7, C_TEXT, 1, cv2.LINE_AA) + cv2.putText(img, 'RATE', (380, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) + cv2.putText(img, f'{rate:.1f}/min', (380, 44), cv2.FONT_HERSHEY_SIMPLEX, 0.7, C_ACCENT, 1, cv2.LINE_AA) + cv2.putText(img, clock, (w - 180, 36), cv2.FONT_HERSHEY_SIMPLEX, 0.55, C_TEXT, 1, cv2.LINE_AA) + cv2.putText(img, f'CAM {camera_id}', (w - 180, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) + + +def draw_footer(img, w, h, frame_idx, live_tag): + bar_h = 28 + overlay_rect(img, 0, h - bar_h, w, h, C_PANEL, alpha=0.55) + cv2.putText(img, f'{live_tag} | Frame {frame_idx}', (12, h - 9), cv2.FONT_HERSHEY_SIMPLEX, 0.45, C_MUTED, 1, cv2.LINE_AA) + + +def draw_skeleton_bold(img, kpts): + for (a, b), color in zip(SKELETON, SK_COLORS): + if a < len(kpts) and b < len(kpts): + xa, ya = int(kpts[a][0]), int(kpts[a][1]) + xb, yb = int(kpts[b][0]), int(kpts[b][1]) + if xa > 0 and ya > 0 and xb > 0 and yb > 0: + cv2.line(img, (xa, ya), (xb, yb), color, 3, cv2.LINE_AA) + for kp in kpts: + x, y = int(kp[0]), int(kp[1]) + if x > 0 and y > 0: + cv2.circle(img, (x, y), 6, (255, 255, 255), -1, cv2.LINE_AA) + cv2.circle(img, (x, y), 6, (40, 40, 40), 2, cv2.LINE_AA) + + +def draw_popups(img, popups, frame_idx): + alive = [] + for pop in popups: + age = frame_idx - pop['born'] + if age > POPUP_LIFETIME: + continue + alive.append(pop) + fade = 1.0 - age / POPUP_LIFETIME + y = pop['y'] - int(age * 1.8) + color = (int(C_GREEN[0] * fade), int(C_GREEN[1] * fade), int(C_GREEN[2] * fade)) + cv2.putText(img, pop['text'], (pop['x'], y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2, cv2.LINE_AA) + return alive + + +def draw_batch_banner(img, w, batch_num, pulse_remaining): + if pulse_remaining <= 0: + return + text = f'NEW BATCH {batch_num}' + font = cv2.FONT_HERSHEY_SIMPLEX + (tw, th), _ = cv2.getTextSize(text, font, 0.8, 2) + x1, y1 = w // 2 - tw // 2 - 16, 62 + x2, y2 = w // 2 + tw // 2 + 16, 62 + th + 20 + overlay_rect(img, x1, y1, x2, y2, C_PANEL, alpha=0.7) + cv2.rectangle(img, (x1, y1), (x2, y2), C_ACCENT, 2) + cv2.putText(img, text, (w // 2 - tw // 2, 62 + th + 4), font, 0.8, C_ACCENT, 2, cv2.LINE_AA) + + +def connect_stream(source, warmup=WARMUP_FRAMES): + attempts = 0 + while not shutdown_requested: + cap = open_capture(source) + if not cap.isOpened(): + attempts += 1 + if MAX_RECONNECT_ATTEMPTS and attempts >= MAX_RECONNECT_ATTEMPTS: + raise RuntimeError(f'Cannot open source after {attempts} attempts: {source}') + print(f'Cannot open source, retry in {RECONNECT_DELAY_SEC}s...') + time.sleep(RECONNECT_DELAY_SEC) + continue + if warmup > 0 and source.lower().startswith(('rtsp://', 'http://')): + warmup_stream(cap, warmup) + w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + fps = cap.get(cv2.CAP_PROP_FPS) + if not fps or fps <= 1: + fps = OUTPUT_FPS + return cap, w, h, fps + return None, 0, 0, OUTPUT_FPS + + +# ============================================================================= +# Main loop +# ============================================================================= + +def run(): + global shutdown_requested + + store = BatchStore( + db_path=DB_PATH, + state_file=STATE_FILE, + camera_name=CAMERA_NAME, + object_label=OBJECT_LABEL, + cutoff_time=DAILY_CUTOFF_TIME, + batch_timeout=BATCH_TIMEOUT_SECONDS, + ignore_batch_label_timeout=IGNORE_BATCH_LABEL_TIMEOUT, + min_object_per_batch=MIN_OBJECT_PER_BATCH, + min_duration_per_batch=MIN_DURATION_PER_BATCH, + logger=lambda msg: print(f'[{now_str()}] {msg}'), + ) + store.start_cutoff_watcher() + + cross_logger = None + if EXPORT_CSV: + cross_logger = CsvLogger(CROSS_CSV, ['batch', 'frame', 'timestamp', 'chicken_id']) + + model = RKNNYOLO( + model_path=MODEL_PATH, + core_mask=CORE_MASK, + imgsz=IMGSZ, + conf=CONF, + num_classes=NUM_CLASSES, + score_sigmoid=SCORE_SIGMOID, + ) + + 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 = 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, + ) + talenan_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, + ) + + 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+ByteTrack counter | {w}x{h} @ {fps}fps | line x={line_x} | cross={CROSS_DIRECTION}') + print(f'Model: {MODEL_PATH} | imgsz={IMGSZ} | core_mask={CORE_MASK}') + print(f'ByteTrack: high_thresh={TRACK_HIGH_THRESH} low_thresh={TRACK_LOW_THRESH} ' + f'match_thresh={TRACK_MATCH_THRESH} buffer={TRACK_BUFFER}') + print(f'DB: {DB_PATH}') + print(f'State: {STATE_FILE}') + + if RECORD_VIDEO: + video_writer = VideoSegmentWriter(OUTPUT_DIR, w, h, fps, VIDEO_SEGMENT_SEC) + + 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 + + detections = model(frame) + + if detections: + ayam_boxes_xyxy = [] + ayam_scores = [] + ayam_kpts_list = [] + ayam_cx_list = [] + + talenan_boxes_xyxy = [] + talenan_scores = [] + talenan_kpts_list = [] + talenan_cx_list = [] + + for det in detections: + bbox = det['bbox'] + score = det['score'] + cls_id = det['cls'] + kpts = det['keypoints'] + cx = (bbox[0] + bbox[2]) / 2.0 + + if cls_id == talenan_cls: + talenan_boxes_xyxy.append(bbox) + talenan_scores.append(score) + talenan_kpts_list.append(kpts) + talenan_cx_list.append(cx) + elif cls_id == ayam_cls: + ayam_boxes_xyxy.append(bbox) + ayam_scores.append(score) + ayam_kpts_list.append(kpts) + ayam_cx_list.append(cx) + + ayam_boxes_xyxy = np.array(ayam_boxes_xyxy, dtype=np.float32) + ayam_scores = np.array(ayam_scores, dtype=np.float32) + talenan_boxes_xyxy = np.array(talenan_boxes_xyxy, dtype=np.float32) + talenan_scores = np.array(talenan_scores, dtype=np.float32) + + ayam_track_map, ayam_det_to_track, ayam_lost_map = ayam_tracker.update( + ayam_boxes_xyxy, ayam_scores) + talenan_track_map, talenan_det_to_track, talenan_lost_map = talenan_tracker.update( + talenan_boxes_xyxy, talenan_scores) + + # Process talenan crossings + for di in range(len(talenan_boxes_xyxy)): + tid = talenan_det_to_track.get(di) + if tid is None: + continue + cx = talenan_cx_list[di] + bbox = talenan_boxes_xyxy[di] + + if tid in talenan_tracked: + prev_cx = talenan_tracked[tid][0] + if crossed_line(prev_cx, cx, line_x) and tid not in talenan_line_crossed: + talenan_line_crossed.add(tid) + if store.record_talenan_crossing(tid): + batch_closed_frame = True + talenan_cross_flash[tid] = CROSS_FLASH_FRAMES + popups.append({ + 'x': int(cx) - 20, + 'y': int((bbox[1] + bbox[3]) / 2), + 'born': frame_idx, + 'text': 'BATCH CLOSED', + }) + talenan_tracked[tid] = (cx, mono) + + # Process ayam crossings (including lost tracks for line-cross continuity) + for di in range(len(ayam_boxes_xyxy)): + tid = ayam_det_to_track.get(di) + if tid is None: + continue + cx = ayam_cx_list[di] + + if tid in ayam_tracked: + prev_cx = ayam_tracked[tid][0] + if crossed_line(prev_cx, cx, line_x) and tid not in ayam_line_crossed: + 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((ayam_boxes_xyxy[di][1] + ayam_boxes_xyxy[di][3]) / 2), + 'born': frame_idx, + 'text': '+1', + }) + ayam_tracked[tid] = (cx, mono) + + # Also track lost tracks for line-crossing continuity + for tid, cx in ayam_lost_map.items(): + if tid not in ayam_tracked: + ayam_tracked[tid] = (cx, mono) + + # Draw talenan + for di in range(len(talenan_boxes_xyxy)): + tid = talenan_det_to_track.get(di) + if tid is None: + continue + bbox = talenan_boxes_xyxy[di] + x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3]) + flash = talenan_cross_flash.get(tid, 0) + color = C_GREEN if flash > 0 else C_TALENAN_BOX + cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3 if flash > 0 else 2) + draw_pill(frame, f'TALENAN {tid}', x1, y1 - 4, color) + + # Draw ayam + for di in range(len(ayam_boxes_xyxy)): + tid = ayam_det_to_track.get(di) + if tid is None: + continue + bbox = ayam_boxes_xyxy[di] + x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3]) + flash = ayam_cross_flash.get(tid, 0) + color = C_GREEN if flash > 0 else C_AYAM_BOX + cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3 if flash > 0 else 2) + draw_pill(frame, f'ID {tid}', x1, y1 - 4, color) + kpts = ayam_kpts_list[di] if di < len(ayam_kpts_list) else None + if kpts is not None: + draw_skeleton_bold(frame, kpts) + + if ayam_crossed_frame: + line_pulse = LINE_PULSE_FRAMES + count_pulse = COUNT_PULSE_FRAMES + if batch_closed_frame: + line_pulse = LINE_PULSE_FRAMES + if batch_started_frame: + batch_pulse = BATCH_PULSE_FRAMES + + batch_num = store.current_batch_number or 0 + batch_count = store.current_batch_count + display_total = store.display_total() + 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-BT' if IS_LIVE else 'FILE-RKNN-BT') + popups = draw_popups(frame, popups, frame_idx) + + for flash_store in (ayam_cross_flash, talenan_cross_flash): + for tid in list(flash_store): + flash_store[tid] -= 1 + if flash_store[tid] <= 0: + del flash_store[tid] + line_pulse = max(0, line_pulse - 1) + count_pulse = max(0, count_pulse - 1) + batch_pulse = max(0, batch_pulse - 1) + + if video_writer is not None: + video_writer.write(frame) + + if LIVE_STREAM_ENABLED and frame_idx % LIVE_STREAM_EVERY_N == 0: + try: + Path(LIVE_STREAM_FRAME_PATH).parent.mkdir(parents=True, exist_ok=True) + _, jpeg = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, LIVE_STREAM_QUALITY]) + with open(LIVE_STREAM_FRAME_PATH, 'wb') as f: + f.write(jpeg.tobytes()) + except Exception: + pass + + frame_idx += 1 + if frame_idx % FLUSH_EVERY_N_FRAMES == 0: + print( + f'[{now_str()}] Frame {frame_idx} | Batch {batch_num}: {batch_count} ' + f'| Total: {display_total} | Uptime {elapsed / 3600:.2f}h' + ) + prune_stale_tracks(ayam_tracked, mono) + prune_stale_tracks(talenan_tracked, mono) + + cap.release() + if video_writer is not None: + video_writer.release() + if cross_logger: + cross_logger.close() + model.release() + store.shutdown() + + print('\n=== Batch Summary (SQLite) ===') + print(f'Database: {DB_PATH}') + + +ayam_tracked = {} +talenan_tracked = {} + + +if __name__ == '__main__': + run() diff --git a/install-services.sh b/install-services.sh new file mode 100644 index 0000000..8a37f94 --- /dev/null +++ b/install-services.sh @@ -0,0 +1,68 @@ +#!/usr/bin/env bash +# Install edge Jetson counter + dashboard; disable legacy MQTT frigate-counter. +# Run on the Jetson: sudo ./install-services.sh +set -euo pipefail + +INSTALL_DIR="${INSTALL_DIR:-/opt/jetson-counter}" +VENV_DIR="${VENV_DIR:-/opt/jetson-counter/venv}" +SERVICE_USER="${SERVICE_USER:-jetson}" + +if [[ "$(id -u)" -ne 0 ]]; then + echo "Run as root: sudo ./install-services.sh" + exit 1 +fi + +if [[ ! -f "${INSTALL_DIR}/.env" ]]; then + echo "Missing ${INSTALL_DIR}/.env" + echo " cp ${INSTALL_DIR}/config.env.example ${INSTALL_DIR}/.env && nano ${INSTALL_DIR}/.env" + exit 1 +fi + +if [[ ! -x "${VENV_DIR}/bin/python" ]]; then + echo "Missing venv: ${VENV_DIR}/bin/python" + echo " sudo ./setup-venv.sh" + exit 1 +fi + +sed -i 's/\r$//' "${INSTALL_DIR}/.env" 2>/dev/null || true + +mkdir -p "${INSTALL_DIR}/.ultralytics" "${INSTALL_DIR}/.torch" +chown -R "${SERVICE_USER}:${SERVICE_USER}" "${INSTALL_DIR}" + +# Disable legacy MQTT counter (replace mode) +for legacy in frigate-counter frigate-counter-dashboard; do + if systemctl is-enabled "${legacy}" &>/dev/null; then + systemctl disable --now "${legacy}" || true + echo "Disabled legacy ${legacy}" + fi +done + +for unit in jetson-counter jetson-counter-dashboard; do + sed -e "s|/opt/jetson-counter|${INSTALL_DIR}|g" \ + -e "s|User=jetson|User=${SERVICE_USER}|g" \ + -e "s|Group=jetson|Group=${SERVICE_USER}|g" \ + "${INSTALL_DIR}/${unit}.service" > "/etc/systemd/system/${unit}.service" + echo "Installed /etc/systemd/system/${unit}.service" +done + +chown -R "${SERVICE_USER}:${SERVICE_USER}" "${INSTALL_DIR}" + +PYTHONNOUSERSITE=1 "${VENV_DIR}/bin/python" -c " +from ultralytics import YOLO +import torch +print('import ok | cuda', torch.cuda.is_available()) +" || { + echo "Import check failed — fix venv before starting services." + exit 1 +} + +systemctl daemon-reload +systemctl reset-failed jetson-counter jetson-counter-dashboard 2>/dev/null || true +systemctl enable jetson-counter jetson-counter-dashboard +systemctl restart jetson-counter jetson-counter-dashboard + +echo "" +systemctl --no-pager status jetson-counter jetson-counter-dashboard || true +echo "" +echo "Logs: sudo journalctl -u jetson-counter -f" +echo "Dashboard: http://$(hostname -I | awk '{print $1}'):5000" diff --git a/jetson-counter-dashboard.service b/jetson-counter-dashboard.service new file mode 100644 index 0000000..7e25d03 --- /dev/null +++ b/jetson-counter-dashboard.service @@ -0,0 +1,33 @@ +[Unit] +Description=Jetson Edge Counter Dashboard (Flask, port 5000) +Documentation=file:///opt/jetson-counter/DEPLOY.md +After=network-online.target jetson-counter.service +Wants=network-online.target + +[Service] +Type=simple +User=jetson +Group=jetson + +WorkingDirectory=/opt/jetson-counter + +EnvironmentFile=/opt/jetson-counter/.env +Environment=PATH=/opt/jetson-counter/venv/bin:/usr/local/bin:/usr/bin:/bin +Environment=FLASK_DEBUG=false + +ExecStart=/opt/jetson-counter/venv/bin/python counter_dashboard.py + +TimeoutStopSec=15 +KillSignal=SIGTERM + +Restart=on-failure +RestartSec=5 +StartLimitInterval=60s +StartLimitBurst=3 + +NoNewPrivileges=true +ProtectHome=true +PrivateTmp=true + +[Install] +WantedBy=multi-user.target diff --git a/jetson-counter.service b/jetson-counter.service new file mode 100644 index 0000000..c24ccb5 --- /dev/null +++ b/jetson-counter.service @@ -0,0 +1,35 @@ +[Unit] +Description=Jetson Edge YOLO Batch Counter (RTSP + TensorRT) +Documentation=file:///opt/jetson-counter/DEPLOY.md +After=network-online.target +Wants=network-online.target + +[Service] +Type=simple +User=jetson +Group=jetson + +WorkingDirectory=/opt/jetson-counter + +EnvironmentFile=/opt/jetson-counter/.env +Environment=PYTHONNOUSERSITE=1 +Environment=YOLO_CONFIG_DIR=/opt/jetson-counter/.ultralytics +Environment=TORCH_HOME=/opt/jetson-counter/.torch +Environment=PATH=/opt/jetson-counter/venv/bin:/usr/local/bin:/usr/bin:/bin + +ExecStart=/opt/jetson-counter/venv/bin/python counter_live.py + +TimeoutStopSec=30 +KillSignal=SIGTERM + +Restart=on-failure +RestartSec=10 +StartLimitInterval=120s +StartLimitBurst=5 + +NoNewPrivileges=true +ProtectHome=true +PrivateTmp=true + +[Install] +WantedBy=multi-user.target diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..fb24f8d --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +numpy<2 +rknn-toolkit-lite2 +opencv-python +flask +python-dotenv diff --git a/setup-venv.sh b/setup-venv.sh new file mode 100644 index 0000000..07116e2 --- /dev/null +++ b/setup-venv.sh @@ -0,0 +1,43 @@ +#!/usr/bin/env bash +# One-time venv for edge Jetson counter — NVIDIA torch required (not PyPI). +set -euo pipefail + +INSTALL_DIR="${INSTALL_DIR:-/opt/jetson-counter}" +VENV_DIR="${VENV_DIR:-/opt/jetson-counter/venv}" +SERVICE_USER="${SERVICE_USER:-jetson}" +TORCH_WHEEL_URL="${TORCH_WHEEL_URL:-https://developer.download.nvidia.com/compute/redist/jp/v60/pytorch/torch-2.4.0a0+3bcc3cddb5.nv24.07.16234504-cp310-cp310-linux_aarch64.whl}" + +if [[ "$(id -u)" -ne 0 ]]; then + echo "Run as root: sudo ./setup-venv.sh" + exit 1 +fi + +apt-get install -y libopenblas-base libopenmpi-dev libomp-dev 2>/dev/null || true + +mkdir -p "${INSTALL_DIR}" +if [[ ! -x "${VENV_DIR}/bin/python" ]]; then + python3 -m venv --system-site-packages "${VENV_DIR}" +fi +chown -R "${SERVICE_USER}:${SERVICE_USER}" "${INSTALL_DIR}" + +sudo -u "${SERVICE_USER}" "${VENV_DIR}/bin/pip" install --upgrade pip +sudo -u "${SERVICE_USER}" "${VENV_DIR}/bin/pip" install "numpy<2" + +if ! sudo -u "${SERVICE_USER}" PYTHONNOUSERSITE=1 "${VENV_DIR}/bin/python" -c "import torch; assert torch.cuda.is_available()" 2>/dev/null; then + echo "Installing NVIDIA Jetson torch wheel..." + sudo -u "${SERVICE_USER}" "${VENV_DIR}/bin/pip" install --no-cache-dir "${TORCH_WHEEL_URL}" +fi + +sudo -u "${SERVICE_USER}" "${VENV_DIR}/bin/pip" install ultralytics flask opencv-python + +sudo -u "${SERVICE_USER}" PYTHONNOUSERSITE=1 "${VENV_DIR}/bin/python" -c " +import torch +from ultralytics import YOLO +import cv2 +import flask +print('venv ok | torch', torch.__version__, '| cuda', torch.cuda.is_available()) +" + +echo "" +echo "If torchvision import fails for ultralytics, copy/build torchvision into venv." +echo "Next: cp config.env.example .env && nano .env && sudo ./install-services.sh" diff --git a/templates/dashboard.html b/templates/dashboard.html new file mode 100644 index 0000000..db1b5ee --- /dev/null +++ b/templates/dashboard.html @@ -0,0 +1,1193 @@ + + + + + +ZenAI APC Dashboard + + + + + + +
+
+
+
+ +
+ +
+
+
◆
+
+

ZenAI APC Module

+
Edge Jetson counter — YOLO TensorRT
+
+
+
+
+ -- + Counting Day +
+
+ + Live +
+ +
+
+ + +
+
+ LIVE +
+
--
+
Current Batch Count
+
Batch #-- · --
+
+
+ Today +
+
--
+
Total Count Today
+
-- batches
+
+
+ Yesterday +
+
--
+
Total Count Yesterday
+
-- batches
+
+
+ Average +
+
--
+
Average Per Day
+
-- days recorded
+
+
+ + +
+ -- + -- +
+ + +
+
+ ◉ Live Camera Feed +
+ + + waiting... + + +
+
+
+
+
▦
+
Live Feed Unavailable
+
Enable LIVE_STREAM_ENABLED=true in counter .env
+
+ +
Connection lost — retrying...
+
+
+ + +
+
+
+ ◉ Daily Trends +
+ + + +
+
+
+ +
+
+
+ ◉ Quick Stats +
+
+
+
+
+
Grand Total
+
All time
+
+
+ -- +
+
+
+
+
+
Total Batches
+
All time
+
+
+ -- +
+
+
+
+
+
Avg per Batch
+
Overall average
+
+
+ -- +
+
+ +
+
◉ Recent Activity
+
+ +
+
+
+
+ + +
+
+
+ ◉ Daily Records +
Click on a row to view batch details
+
+
+ + + + +
+
+
+ + + + + + + + + + + + + + +
DateTotal CountBatchesAvg/BatchStatusAction
+
+
+ +
Jetson Edge Counter — port 5000
+
+ + + + + + + + + + diff --git a/uninstall-services.sh b/uninstall-services.sh new file mode 100644 index 0000000..1f3fa11 --- /dev/null +++ b/uninstall-services.sh @@ -0,0 +1,18 @@ +#!/usr/bin/env bash +# Remove Jetson edge counter systemd services. +# Run on the Jetson: sudo ./uninstall-services.sh +set -euo pipefail + +if [[ "$(id -u)" -ne 0 ]]; then + echo "Run as root: sudo ./uninstall-services.sh" + exit 1 +fi + +for unit in jetson-counter jetson-counter-dashboard; do + systemctl stop "${unit}" 2>/dev/null || true + systemctl disable "${unit}" 2>/dev/null || true + rm -f "/etc/systemd/system/${unit}.service" +done + +systemctl daemon-reload +echo "Removed jetson-counter and jetson-counter-dashboard services."