forked from dsutanto/bytetrack-counter-cpp
First Commit
This commit is contained in:
proitlab committed
1 parent
267107ff05
commit
37fe4eee3a
53 files changed
+15098
-2
No files matched your search
+379
@@ -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()
|
||||
Reference in new issue
Block a user