forked from dsutanto/bytetrack-counter-cpp
380 lines
14 KiB
Python
380 lines
14 KiB
Python
"""
|
||
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()
|