466 lines
17 KiB
Python
466 lines
17 KiB
Python
"""
|
|
Production daily counter persistence for edge sack zone counter.
|
|
Tracks per-zone entries and the daily total per counting day,
|
|
delimited by the daily cutoff time. SQLite schema + current_counter.json state.
|
|
|
|
Zone ids are stable keys like zone_1, zone_2, … (from zones_config).
|
|
"""
|
|
import json
|
|
import sqlite3
|
|
import threading
|
|
import time
|
|
from datetime import datetime, timedelta
|
|
from pathlib import Path
|
|
|
|
|
|
class CounterStore:
|
|
def __init__(
|
|
self,
|
|
db_path,
|
|
state_file,
|
|
camera_name,
|
|
object_label='object',
|
|
cutoff_time='20:00',
|
|
carry_ids=50,
|
|
zone_ids=None,
|
|
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.carry_ids = int(carry_ids)
|
|
self.zone_ids = list(zone_ids) if zone_ids else ['zone_1']
|
|
if not self.zone_ids:
|
|
raise ValueError('zone_ids must contain at least one zone id')
|
|
self._zone_id_set = set(self.zone_ids)
|
|
self.log = logger
|
|
|
|
self.state_lock = threading.Lock()
|
|
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()
|
|
with self.state_lock:
|
|
self._fill_missing_days()
|
|
counting_date = self.get_counting_date()
|
|
if self.current_state is None:
|
|
self._start_new_day(counting_date)
|
|
self._persist_day()
|
|
|
|
def _table_columns(self, table):
|
|
cur = self.db.cursor()
|
|
cur.execute(f'PRAGMA table_info({table})')
|
|
return {row[1] for row in cur.fetchall()}
|
|
|
|
def _init_db(self):
|
|
cur = self.db.cursor()
|
|
cur.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS daily_counters (
|
|
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,
|
|
start_time TEXT,
|
|
end_time TEXT,
|
|
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
UNIQUE(counting_date, camera_name, object_label)
|
|
)
|
|
"""
|
|
)
|
|
cur.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS daily_zone_counts (
|
|
counting_date TEXT NOT NULL,
|
|
camera_name TEXT NOT NULL,
|
|
object_label TEXT NOT NULL,
|
|
zone_id TEXT NOT NULL,
|
|
count INTEGER NOT NULL DEFAULT 0,
|
|
PRIMARY KEY (counting_date, camera_name, object_label, zone_id)
|
|
)
|
|
"""
|
|
)
|
|
self.db.commit()
|
|
self._migrate_left_right_if_needed()
|
|
|
|
def _migrate_left_right_if_needed(self):
|
|
cols = self._table_columns('daily_counters')
|
|
if 'total_left' not in cols and 'total_right' not in cols:
|
|
return
|
|
|
|
self.log('Migrating daily_counters left/right columns to daily_zone_counts')
|
|
cur = self.db.cursor()
|
|
cur.execute(
|
|
"""
|
|
SELECT counting_date, camera_name, object_label,
|
|
COALESCE(total_left, 0), COALESCE(total_right, 0)
|
|
FROM daily_counters
|
|
"""
|
|
)
|
|
rows = cur.fetchall()
|
|
for counting_date, camera_name, object_label, left, right in rows:
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO daily_zone_counts
|
|
(counting_date, camera_name, object_label, zone_id, count)
|
|
VALUES (?, ?, ?, 'zone_1', ?)
|
|
ON CONFLICT(counting_date, camera_name, object_label, zone_id)
|
|
DO UPDATE SET count = excluded.count
|
|
""",
|
|
(counting_date, camera_name, object_label, left),
|
|
)
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO daily_zone_counts
|
|
(counting_date, camera_name, object_label, zone_id, count)
|
|
VALUES (?, ?, ?, 'zone_2', ?)
|
|
ON CONFLICT(counting_date, camera_name, object_label, zone_id)
|
|
DO UPDATE SET count = excluded.count
|
|
""",
|
|
(counting_date, camera_name, object_label, right),
|
|
)
|
|
|
|
cur.execute(
|
|
"""
|
|
CREATE TABLE daily_counters_new (
|
|
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,
|
|
start_time TEXT,
|
|
end_time TEXT,
|
|
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
UNIQUE(counting_date, camera_name, object_label)
|
|
)
|
|
"""
|
|
)
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO daily_counters_new
|
|
(id, counting_date, camera_name, object_label,
|
|
total_count, start_time, end_time, updated_at)
|
|
SELECT id, counting_date, camera_name, object_label,
|
|
total_count, start_time, end_time, updated_at
|
|
FROM daily_counters
|
|
"""
|
|
)
|
|
cur.execute('DROP TABLE daily_counters')
|
|
cur.execute('ALTER TABLE daily_counters_new RENAME TO daily_counters')
|
|
self.db.commit()
|
|
self.log(f'Migration complete ({len(rows)} daily row(s))')
|
|
|
|
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 _empty_counts(self):
|
|
return {zid: 0 for zid in self.zone_ids}
|
|
|
|
def _normalize_state(self, state):
|
|
counts = state.get('counts')
|
|
if not isinstance(counts, dict):
|
|
counts = {}
|
|
if 'count_left' in state or 'count_right' in state:
|
|
counts['zone_1'] = int(state.get('count_left', 0) or 0)
|
|
counts['zone_2'] = int(state.get('count_right', 0) or 0)
|
|
normalized = self._empty_counts()
|
|
for zid in self.zone_ids:
|
|
normalized[zid] = int(counts.get(zid, 0) or 0)
|
|
# Preserve migrated zone_2 even if current config is 1 zone? Spec says
|
|
# configured ids only — drop unknown keys from active counts.
|
|
state['counts'] = normalized
|
|
state['total_count'] = sum(normalized.values())
|
|
state.pop('count', None)
|
|
state.pop('count_left', None)
|
|
state.pop('count_right', None)
|
|
state.setdefault('counted_event_ids', [])
|
|
state.setdefault('start_time', datetime.now().isoformat())
|
|
state.setdefault('last_detection_time', state.get('start_time'))
|
|
return state
|
|
|
|
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 "
|
|
f"({state.get('counting_date')}). Starting fresh."
|
|
)
|
|
self.state_file.unlink(missing_ok=True)
|
|
return None
|
|
state = self._normalize_state(state)
|
|
stale_ids = state.get('counted_event_ids') or []
|
|
if stale_ids:
|
|
self.log(
|
|
f"Cleared {len(stale_ids)} persisted track dedup keys "
|
|
f"(track IDs reset on restart)"
|
|
)
|
|
state['counted_event_ids'] = []
|
|
self.log(
|
|
f"Resumed {current_date} with total={state['total_count']} "
|
|
f"counts={state['counts']}"
|
|
)
|
|
with open(self.state_file, 'w', encoding='utf-8') as f:
|
|
json.dump(state, f, indent=2, ensure_ascii=False)
|
|
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 _start_new_day(self, counting_date):
|
|
now = datetime.now().isoformat()
|
|
carried = []
|
|
if self.current_state is not None:
|
|
try:
|
|
carried = self.current_state['counted_event_ids'][-self.carry_ids:]
|
|
except (KeyError, TypeError):
|
|
carried = []
|
|
counts = self._empty_counts()
|
|
self.current_state = {
|
|
'counting_date': counting_date,
|
|
'total_count': 0,
|
|
'counts': counts,
|
|
'start_time': now,
|
|
'last_detection_time': now,
|
|
'counted_event_ids': carried,
|
|
}
|
|
self.save_state()
|
|
self.log(f'Started counting day {counting_date} ({self.object_label})')
|
|
|
|
def _ensure_day_row(self, counting_date, start_time=None, end_time=None):
|
|
"""Insert a zero row for counting_date if it does not already exist."""
|
|
now = datetime.now().isoformat()
|
|
cur = self.db.cursor()
|
|
cur.execute(
|
|
"""
|
|
INSERT OR IGNORE INTO daily_counters
|
|
(counting_date, camera_name, object_label,
|
|
total_count, start_time, end_time)
|
|
VALUES (?, ?, ?, 0, ?, ?)
|
|
""",
|
|
(
|
|
counting_date,
|
|
self.camera_name,
|
|
self.object_label,
|
|
start_time or now,
|
|
end_time or now,
|
|
),
|
|
)
|
|
for zid in self.zone_ids:
|
|
cur.execute(
|
|
"""
|
|
INSERT OR IGNORE INTO daily_zone_counts
|
|
(counting_date, camera_name, object_label, zone_id, count)
|
|
VALUES (?, ?, ?, ?, 0)
|
|
""",
|
|
(counting_date, self.camera_name, self.object_label, zid),
|
|
)
|
|
self.db.commit()
|
|
|
|
def _fill_missing_days(self):
|
|
"""Backfill any missing counting dates from first DB row through today as 0."""
|
|
today = self.get_counting_date()
|
|
cur = self.db.cursor()
|
|
cur.execute(
|
|
"""
|
|
SELECT counting_date FROM daily_counters
|
|
WHERE camera_name = ? AND object_label = ?
|
|
ORDER BY counting_date ASC
|
|
""",
|
|
(self.camera_name, self.object_label),
|
|
)
|
|
existing = {row[0] for row in cur.fetchall()}
|
|
if not existing:
|
|
self._ensure_day_row(today)
|
|
return
|
|
|
|
start = datetime.strptime(min(existing), '%Y-%m-%d').date()
|
|
end = datetime.strptime(today, '%Y-%m-%d').date()
|
|
filled = 0
|
|
d = start
|
|
while d <= end:
|
|
key = d.isoformat()
|
|
if key not in existing:
|
|
self._ensure_day_row(key)
|
|
filled += 1
|
|
d += timedelta(days=1)
|
|
if filled:
|
|
self.log(
|
|
f'Backfilled {filled} zero-activity day(s) '
|
|
f'{start.isoformat()}..{end.isoformat()}'
|
|
)
|
|
|
|
def record_zone_entry(self, track_id, zone_id):
|
|
"""Record a sack entering a zone. zone_id: 'zone_1' | 'zone_2' | …
|
|
|
|
Returns (total_count, day_started, counted).
|
|
"""
|
|
if zone_id not in self._zone_id_set:
|
|
raise ValueError(
|
|
f"zone_id must be one of {self.zone_ids}, got {zone_id!r}"
|
|
)
|
|
with self.state_lock:
|
|
counting_date = self.get_counting_date()
|
|
day_started = False
|
|
if self.current_state is None or self.current_state['counting_date'] != counting_date:
|
|
self._start_new_day(counting_date)
|
|
day_started = True
|
|
|
|
counted = False
|
|
event_key = f"{track_id}_{zone_id}"
|
|
if event_key not in self.current_state['counted_event_ids']:
|
|
self.current_state['total_count'] += 1
|
|
self.current_state['counts'][zone_id] = (
|
|
int(self.current_state['counts'].get(zone_id, 0)) + 1
|
|
)
|
|
self.current_state['counted_event_ids'].append(event_key)
|
|
counted = True
|
|
self.log(
|
|
f'Counted {zone_id} (track {track_id}) | {counting_date} '
|
|
f'total: {self.current_state["total_count"]} '
|
|
f'counts={self.current_state["counts"]}'
|
|
)
|
|
self._persist_day()
|
|
|
|
self.current_state['last_detection_time'] = datetime.now().isoformat()
|
|
self.save_state()
|
|
return self.current_state['total_count'], day_started, counted
|
|
|
|
def _persist_day(self):
|
|
state = self.current_state
|
|
cur = self.db.cursor()
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO daily_counters
|
|
(counting_date, camera_name, object_label,
|
|
total_count, start_time, end_time)
|
|
VALUES (?, ?, ?, ?, ?, ?)
|
|
ON CONFLICT(counting_date, camera_name, object_label)
|
|
DO UPDATE SET
|
|
total_count = excluded.total_count,
|
|
end_time = excluded.end_time,
|
|
updated_at = CURRENT_TIMESTAMP
|
|
""",
|
|
(
|
|
state['counting_date'], self.camera_name, self.object_label,
|
|
state['total_count'],
|
|
state['start_time'], datetime.now().isoformat(),
|
|
),
|
|
)
|
|
for zid, count in state['counts'].items():
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO daily_zone_counts
|
|
(counting_date, camera_name, object_label, zone_id, count)
|
|
VALUES (?, ?, ?, ?, ?)
|
|
ON CONFLICT(counting_date, camera_name, object_label, zone_id)
|
|
DO UPDATE SET count = excluded.count
|
|
""",
|
|
(
|
|
state['counting_date'], self.camera_name, self.object_label,
|
|
zid, int(count),
|
|
),
|
|
)
|
|
self.db.commit()
|
|
|
|
def cutoff_watcher_loop(self):
|
|
while not self.shutdown_event.is_set():
|
|
time.sleep(60)
|
|
with self.state_lock:
|
|
self._fill_missing_days()
|
|
counting_date = self.get_counting_date()
|
|
if self.current_state is None:
|
|
self._start_new_day(counting_date)
|
|
self._persist_day()
|
|
continue
|
|
if self.current_state['counting_date'] != counting_date:
|
|
self.log('Daily cutoff reached - finalizing day totals')
|
|
self._persist_day()
|
|
self.current_state = None
|
|
self.save_state()
|
|
self._start_new_day(counting_date)
|
|
self._persist_day()
|
|
|
|
def start_cutoff_watcher(self):
|
|
t = threading.Thread(target=self.cutoff_watcher_loop, daemon=True)
|
|
t.start()
|
|
return t
|
|
|
|
@property
|
|
def current_count(self):
|
|
if self.current_state is None:
|
|
return 0
|
|
return self.current_state['total_count']
|
|
|
|
@property
|
|
def current_counts(self):
|
|
if self.current_state is None:
|
|
return self._empty_counts()
|
|
return dict(self.current_state.get('counts') or self._empty_counts())
|
|
|
|
def _day_zone_counts(self, counting_date=None):
|
|
if counting_date is None:
|
|
counting_date = self.get_counting_date()
|
|
cur = self.db.cursor()
|
|
cur.execute(
|
|
"""
|
|
SELECT zone_id, COALESCE(count, 0)
|
|
FROM daily_zone_counts
|
|
WHERE counting_date = ? AND camera_name = ? AND object_label = ?
|
|
""",
|
|
(counting_date, self.camera_name, self.object_label),
|
|
)
|
|
out = self._empty_counts()
|
|
for zone_id, count in cur.fetchall():
|
|
if zone_id in out:
|
|
out[zone_id] = int(count)
|
|
return out
|
|
|
|
def display_total(self):
|
|
cur = self.db.cursor()
|
|
cur.execute(
|
|
"""
|
|
SELECT COALESCE(total_count, 0)
|
|
FROM daily_counters
|
|
WHERE counting_date = ? AND camera_name = ? AND object_label = ?
|
|
""",
|
|
(self.get_counting_date(), self.camera_name, self.object_label),
|
|
)
|
|
row = cur.fetchone()
|
|
return int(row[0]) if row else 0
|
|
|
|
def display_counts(self):
|
|
return self._day_zone_counts()
|
|
|
|
def shutdown(self):
|
|
self.shutdown_event.set()
|
|
with self.state_lock:
|
|
if self.current_state is not None:
|
|
self._persist_day()
|
|
self.db.close()
|