482 lines
20 KiB
Python
482 lines
20 KiB
Python
"""SQLite storage for metadata and status.
|
|
|
|
The split is deliberate: this database holds *what* and *where*, the disk holds
|
|
the pixels, the final YOLO labels, and the weights. A master dataset stays
|
|
trainable even if this file is deleted (REQ-006, REQ-054).
|
|
|
|
Connections are per-call rather than shared, because the job worker runs on its
|
|
own thread and SQLite connections are not safely shared across threads. WAL mode
|
|
lets that worker write while requests read.
|
|
"""
|
|
|
|
import os
|
|
import sqlite3
|
|
from contextlib import contextmanager
|
|
from typing import Iterator
|
|
|
|
from backend import config
|
|
|
|
SCHEMA = [
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS projects (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
slug TEXT NOT NULL UNIQUE,
|
|
name TEXT NOT NULL,
|
|
label_type TEXT NOT NULL CHECK (label_type IN ('bbox', 'polygon')),
|
|
base_model_path TEXT,
|
|
base_model_kind TEXT CHECK (base_model_kind IN ('uploaded', 'pretrained', 'trained')),
|
|
video_root TEXT NOT NULL,
|
|
val_every INTEGER NOT NULL DEFAULT 5,
|
|
created_at REAL NOT NULL
|
|
)
|
|
""",
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS project_classes (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
project_id INTEGER NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
|
|
class_id INTEGER NOT NULL,
|
|
name TEXT NOT NULL,
|
|
prompt TEXT NOT NULL,
|
|
UNIQUE (project_id, class_id)
|
|
)
|
|
""",
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS batches (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
project_id INTEGER NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
|
|
video_path TEXT NOT NULL,
|
|
date_label TEXT NOT NULL,
|
|
batch_label TEXT NOT NULL,
|
|
start_sec REAL NOT NULL,
|
|
end_sec REAL NOT NULL,
|
|
fps REAL NOT NULL,
|
|
status TEXT NOT NULL CHECK (status IN (
|
|
'extracting', 'extracted', 'labeling', 'reviewing',
|
|
'approved', 'merged', 'failed')),
|
|
frame_count INTEGER NOT NULL DEFAULT 0,
|
|
created_at REAL NOT NULL,
|
|
merged_at REAL
|
|
)
|
|
""",
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS frames (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
batch_id INTEGER NOT NULL REFERENCES batches(id) ON DELETE CASCADE,
|
|
idx INTEGER NOT NULL,
|
|
filename TEXT NOT NULL,
|
|
width INTEGER NOT NULL,
|
|
height INTEGER NOT NULL,
|
|
review_status TEXT NOT NULL DEFAULT 'pending'
|
|
CHECK (review_status IN ('pending', 'approved', 'rejected')),
|
|
UNIQUE (batch_id, idx)
|
|
)
|
|
""",
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS annotations (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
frame_id INTEGER NOT NULL REFERENCES frames(id) ON DELETE CASCADE,
|
|
class_id INTEGER NOT NULL,
|
|
geometry TEXT NOT NULL,
|
|
score REAL NOT NULL DEFAULT 1.0,
|
|
source TEXT NOT NULL CHECK (source IN ('auto', 'manual')),
|
|
created_at REAL NOT NULL
|
|
)
|
|
""",
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS datasets (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
project_id INTEGER NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
|
|
name TEXT NOT NULL,
|
|
note TEXT NOT NULL DEFAULT '',
|
|
rule_version TEXT,
|
|
rules_json TEXT,
|
|
created_at REAL NOT NULL
|
|
)
|
|
""",
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS dataset_items (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
project_id INTEGER NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
|
|
dataset_id INTEGER NOT NULL REFERENCES datasets(id) ON DELETE CASCADE,
|
|
frame_id INTEGER NOT NULL REFERENCES frames(id) ON DELETE CASCADE,
|
|
split TEXT NOT NULL CHECK (split IN ('train', 'val')),
|
|
image_rel TEXT NOT NULL,
|
|
label_rel TEXT NOT NULL,
|
|
added_at REAL NOT NULL,
|
|
UNIQUE (dataset_id, frame_id)
|
|
)
|
|
""",
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS base_datasets (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
project_id INTEGER NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
|
|
name TEXT NOT NULL,
|
|
source TEXT NOT NULL DEFAULT '',
|
|
image_count INTEGER NOT NULL DEFAULT 0,
|
|
box_count INTEGER NOT NULL DEFAULT 0,
|
|
classes TEXT NOT NULL DEFAULT '[]',
|
|
created_at REAL NOT NULL
|
|
)
|
|
""",
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS video_clock (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
project_id INTEGER NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
|
|
video_rel TEXT NOT NULL,
|
|
folder_date TEXT NOT NULL DEFAULT '',
|
|
-- Wall-clock text exactly as the camera burned it in, e.g.
|
|
-- '2026-08-06 20:12:42'. Deliberately not an epoch: the overlay carries
|
|
-- no timezone, and storing it as one made the backend (UTC) and the
|
|
-- browser (UTC+7) disagree by 7 hours, which moved recordings across
|
|
-- the 06:00 boundary into the wrong working day.
|
|
started_at TEXT,
|
|
working_day TEXT NOT NULL DEFAULT '',
|
|
confidence REAL NOT NULL DEFAULT 0,
|
|
agreeing INTEGER NOT NULL DEFAULT 0,
|
|
source TEXT NOT NULL DEFAULT 'ocr',
|
|
error TEXT NOT NULL DEFAULT '',
|
|
read_at REAL,
|
|
UNIQUE (project_id, video_rel)
|
|
)
|
|
""",
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS count_runs (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
project_id INTEGER NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
|
|
video_rel TEXT NOT NULL,
|
|
date_label TEXT NOT NULL,
|
|
batch_label TEXT NOT NULL,
|
|
loading INTEGER,
|
|
unloading INTEGER,
|
|
net INTEGER,
|
|
ground_truth INTEGER,
|
|
frames INTEGER NOT NULL DEFAULT 0,
|
|
seconds REAL NOT NULL DEFAULT 0,
|
|
params TEXT NOT NULL DEFAULT '{}',
|
|
model_path TEXT NOT NULL DEFAULT '',
|
|
error TEXT NOT NULL DEFAULT '',
|
|
counted_at REAL,
|
|
UNIQUE (project_id, video_rel)
|
|
)
|
|
""",
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS model_versions (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
project_id INTEGER NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
|
|
version INTEGER NOT NULL,
|
|
weights_path TEXT NOT NULL,
|
|
parent_model_path TEXT,
|
|
metrics TEXT,
|
|
base_metrics TEXT,
|
|
created_at REAL NOT NULL,
|
|
UNIQUE (project_id, version)
|
|
)
|
|
""",
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS jobs (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
project_id INTEGER REFERENCES projects(id) ON DELETE CASCADE,
|
|
batch_id INTEGER REFERENCES batches(id) ON DELETE CASCADE,
|
|
type TEXT NOT NULL CHECK (type IN ('extract', 'autolabel', 'merge', 'train',
|
|
'count', 'clock-scan',
|
|
'truck-scan')),
|
|
status TEXT NOT NULL CHECK (status IN (
|
|
'queued', 'running', 'done', 'failed', 'cancelled')),
|
|
params TEXT NOT NULL DEFAULT '{}',
|
|
progress INTEGER NOT NULL DEFAULT 0,
|
|
total INTEGER NOT NULL DEFAULT 0,
|
|
message TEXT NOT NULL DEFAULT '',
|
|
error TEXT NOT NULL DEFAULT '',
|
|
log TEXT NOT NULL DEFAULT '',
|
|
created_at REAL NOT NULL,
|
|
started_at REAL,
|
|
finished_at REAL
|
|
)
|
|
""",
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS triage_rules (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
project_id INTEGER NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
|
|
stage TEXT NOT NULL DEFAULT 'dataprep',
|
|
position INTEGER NOT NULL,
|
|
name TEXT NOT NULL,
|
|
predicate TEXT NOT NULL,
|
|
action TEXT NOT NULL CHECK (action IN ('keep', 'ignore', 'reclass')),
|
|
target_class INTEGER,
|
|
created_at REAL NOT NULL
|
|
)
|
|
""",
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS annotation_overrides (
|
|
annotation_id INTEGER PRIMARY KEY REFERENCES annotations(id) ON DELETE CASCADE,
|
|
verdict TEXT NOT NULL CHECK (verdict IN ('keep', 'ignore', 'reclass')),
|
|
target_class INTEGER,
|
|
decided_at REAL NOT NULL
|
|
)
|
|
""",
|
|
"CREATE INDEX IF NOT EXISTS idx_frames_batch ON frames(batch_id, idx)",
|
|
"CREATE INDEX IF NOT EXISTS idx_triage_rules_project ON triage_rules(project_id, stage, position)",
|
|
"CREATE INDEX IF NOT EXISTS idx_annotations_frame ON annotations(frame_id)",
|
|
"CREATE INDEX IF NOT EXISTS idx_batches_project ON batches(project_id)",
|
|
"CREATE INDEX IF NOT EXISTS idx_jobs_project ON jobs(project_id, created_at)",
|
|
"CREATE INDEX IF NOT EXISTS idx_dataset_items_project ON dataset_items(project_id)",
|
|
"CREATE INDEX IF NOT EXISTS idx_count_runs_project ON count_runs(project_id, date_label)",
|
|
"CREATE INDEX IF NOT EXISTS idx_video_clock_project ON video_clock(project_id, working_day)",
|
|
]
|
|
|
|
|
|
def connect() -> sqlite3.Connection:
|
|
os.makedirs(os.path.dirname(config.DB_PATH), exist_ok=True)
|
|
connection = sqlite3.connect(config.DB_PATH, timeout=30.0)
|
|
connection.row_factory = sqlite3.Row
|
|
connection.execute("PRAGMA journal_mode = WAL")
|
|
connection.execute("PRAGMA foreign_keys = ON")
|
|
connection.execute("PRAGMA busy_timeout = 30000")
|
|
return connection
|
|
|
|
|
|
@contextmanager
|
|
def cursor() -> Iterator[sqlite3.Cursor]:
|
|
"""Transactional cursor: commits on success, rolls back on exception."""
|
|
connection = connect()
|
|
try:
|
|
with connection:
|
|
yield connection.cursor()
|
|
finally:
|
|
connection.close()
|
|
|
|
|
|
def migrate() -> None:
|
|
"""Create every table and index. Idempotent — safe on every startup."""
|
|
with cursor() as cur:
|
|
for statement in SCHEMA:
|
|
cur.execute(statement)
|
|
cur.execute("PRAGMA table_info(projects)")
|
|
cols = [column[1] for column in cur.fetchall()]
|
|
if "secondary_model_path" not in cols:
|
|
cur.execute("ALTER TABLE projects ADD COLUMN secondary_model_path TEXT")
|
|
if "secondary_model_name" not in cols:
|
|
cur.execute("ALTER TABLE projects ADD COLUMN secondary_model_name TEXT")
|
|
if "secondary_model_classes" not in cols:
|
|
cur.execute("ALTER TABLE projects ADD COLUMN secondary_model_classes TEXT")
|
|
# REQ-110: augmentation settings, null until the user changes them.
|
|
if "augment" not in cols:
|
|
cur.execute("ALTER TABLE projects ADD COLUMN augment TEXT")
|
|
# REQ-184: a class flagged container keeps boxes that sit inside it.
|
|
cur.execute("PRAGMA table_info(project_classes)")
|
|
class_cols = [column[1] for column in cur.fetchall()]
|
|
if "container" not in class_cols:
|
|
cur.execute(
|
|
"ALTER TABLE project_classes ADD COLUMN container INTEGER NOT NULL DEFAULT 0"
|
|
)
|
|
# REQ-107: what rule set a run's numbers were measured under.
|
|
cur.execute("PRAGMA table_info(model_versions)")
|
|
version_cols = [column[1] for column in cur.fetchall()]
|
|
if "rule_version" not in version_cols:
|
|
cur.execute("ALTER TABLE model_versions ADD COLUMN rule_version TEXT")
|
|
# REQ-113: and what augmentation it trained under.
|
|
if "augment" not in version_cols:
|
|
cur.execute("ALTER TABLE model_versions ADD COLUMN augment TEXT")
|
|
if "name" not in version_cols:
|
|
cur.execute("ALTER TABLE model_versions ADD COLUMN name TEXT")
|
|
|
|
# REQ-132: the triage rules this dataset was actually cut under, frozen
|
|
# at merge time. Editing project rules afterwards must not rewrite what
|
|
# an existing dataset claims.
|
|
cur.execute("PRAGMA table_info(datasets)")
|
|
dataset_cols = [column[1] for column in cur.fetchall()]
|
|
if "rules_json" not in dataset_cols:
|
|
cur.execute("ALTER TABLE datasets ADD COLUMN rules_json TEXT")
|
|
|
|
_migrate_dataset_items(cur)
|
|
_backfill_dataset_rules(cur)
|
|
_migrate_job_types(cur)
|
|
_migrate_clock_column(cur)
|
|
_migrate_truck_columns(cur)
|
|
_backfill_model_names(cur)
|
|
|
|
|
|
def _migrate_dataset_items(cur) -> None:
|
|
"""Move a single-dataset project onto named datasets.
|
|
|
|
The old table pinned `frame_id UNIQUE`, so one frame could live in exactly
|
|
one dataset — which is the constraint that made merging the same batch twice
|
|
a conflict. SQLite cannot drop a constraint, so the table is rebuilt and the
|
|
existing rows are adopted by a "Master Dataset #1" per project. Nothing is
|
|
discarded: a row without a home would be a merged frame the app can no
|
|
longer see.
|
|
"""
|
|
cur.execute("PRAGMA table_info(dataset_items)")
|
|
columns = [column[1] for column in cur.fetchall()]
|
|
if not columns or "dataset_id" in columns:
|
|
return
|
|
|
|
import time
|
|
|
|
cur.execute("SELECT DISTINCT project_id FROM dataset_items")
|
|
projects_with_items = [row[0] for row in cur.fetchall()]
|
|
home = {}
|
|
for project_id in projects_with_items:
|
|
cur.execute(
|
|
"""INSERT INTO datasets (project_id, name, note, created_at)
|
|
VALUES (?, 'Master Dataset #1', 'Adopted from the original single dataset', ?)""",
|
|
(project_id, time.time()),
|
|
)
|
|
home[project_id] = cur.lastrowid
|
|
|
|
cur.execute("ALTER TABLE dataset_items RENAME TO dataset_items_old")
|
|
cur.execute("""
|
|
CREATE TABLE dataset_items (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
project_id INTEGER NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
|
|
dataset_id INTEGER NOT NULL REFERENCES datasets(id) ON DELETE CASCADE,
|
|
frame_id INTEGER NOT NULL REFERENCES frames(id) ON DELETE CASCADE,
|
|
split TEXT NOT NULL CHECK (split IN ('train', 'val')),
|
|
image_rel TEXT NOT NULL,
|
|
label_rel TEXT NOT NULL,
|
|
added_at REAL NOT NULL,
|
|
UNIQUE (dataset_id, frame_id)
|
|
)
|
|
""")
|
|
for project_id, dataset_id in home.items():
|
|
cur.execute(
|
|
"""INSERT INTO dataset_items (project_id, dataset_id, frame_id, split,
|
|
image_rel, label_rel, added_at)
|
|
SELECT project_id, ?, frame_id, split, image_rel, label_rel, added_at
|
|
FROM dataset_items_old WHERE project_id = ?""",
|
|
(dataset_id, project_id),
|
|
)
|
|
cur.execute("DROP TABLE dataset_items_old")
|
|
|
|
|
|
def _migrate_clock_column(cur) -> None:
|
|
"""Drop a `video_clock` that still stores epochs.
|
|
|
|
The table is a cache of what the reader saw; rebuilding it is one rescan, and
|
|
that is cheaper and safer than converting epochs whose timezone is exactly
|
|
the thing that was wrong.
|
|
"""
|
|
cur.execute("PRAGMA table_info(video_clock)")
|
|
columns = {row[1]: row[2] for row in cur.fetchall()}
|
|
if columns.get("started_at", "TEXT").upper() == "REAL":
|
|
cur.execute("DROP TABLE video_clock")
|
|
for statement in SCHEMA:
|
|
if "CREATE TABLE IF NOT EXISTS video_clock" in statement:
|
|
cur.execute(statement)
|
|
break
|
|
|
|
|
|
def _migrate_truck_columns(cur) -> None:
|
|
"""Truck-presence check per recording (REQ-166)."""
|
|
cur.execute("PRAGMA table_info(video_clock)")
|
|
columns = {row[1] for row in cur.fetchall()}
|
|
for name, kind in (("truck_hits", "INTEGER"), ("truck_samples", "INTEGER"),
|
|
("truck_model", "TEXT"), ("truck_checked_at", "REAL")):
|
|
if name not in columns:
|
|
cur.execute(f"ALTER TABLE video_clock ADD COLUMN {name} {kind}")
|
|
|
|
|
|
def _migrate_job_types(cur) -> None:
|
|
"""Let the jobs table accept the 'count' type (REQ-150).
|
|
|
|
The allowed set is a CHECK constraint written into the table definition, and
|
|
SQLite cannot alter one — the table has to be rebuilt. Rows are carried over
|
|
verbatim; the only change is which type strings are legal.
|
|
"""
|
|
cur.execute("SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'jobs'")
|
|
row = cur.fetchone()
|
|
if row is None or "'truck-scan'" in row[0]:
|
|
return
|
|
cur.execute("ALTER TABLE jobs RENAME TO jobs_old")
|
|
for statement in SCHEMA:
|
|
if "CREATE TABLE IF NOT EXISTS jobs " in statement:
|
|
cur.execute(statement)
|
|
break
|
|
cur.execute("""INSERT INTO jobs (id, project_id, batch_id, type, status, params, progress,
|
|
total, message, error, log, created_at, started_at,
|
|
finished_at)
|
|
SELECT id, project_id, batch_id, type, status, params, progress, total,
|
|
message, error, log, created_at, started_at, finished_at
|
|
FROM jobs_old""")
|
|
cur.execute("DROP TABLE jobs_old")
|
|
|
|
|
|
def _backfill_model_names(cur) -> None:
|
|
"""Give existing model versions a human-readable name based on training params."""
|
|
import time as _time
|
|
|
|
PRETRAINED = {"bbox": "yolo11n.pt", "polygon": "yolo11n-seg.pt"}
|
|
|
|
cur.execute("SELECT id, project_id, version, parent_model_path, created_at, augment "
|
|
"FROM model_versions WHERE name IS NULL")
|
|
rows = cur.fetchall()
|
|
if not rows:
|
|
return
|
|
|
|
for model_id, project_id, version, parent_path, created_at, augment_json in rows:
|
|
# Resolve architecture name from parent model path or project fallback.
|
|
cur.execute("SELECT label_type FROM projects WHERE id = ?", (project_id,))
|
|
proj_row = cur.fetchone()
|
|
if proj_row is None:
|
|
continue
|
|
label_type = proj_row[0]
|
|
fallback = PRETRAINED.get(label_type, "yolo11n.pt")
|
|
|
|
if parent_path:
|
|
arch = os.path.splitext(os.path.basename(parent_path))[0]
|
|
else:
|
|
arch = os.path.splitext(os.path.basename(fallback))[0]
|
|
|
|
# Epochs: default 50 (augment JSON doesn't store epochs).
|
|
epochs = 50
|
|
|
|
# Class names.
|
|
cur.execute(
|
|
"SELECT name FROM project_classes WHERE project_id = ? ORDER BY class_id",
|
|
(project_id,),
|
|
)
|
|
class_names = "+".join(r[0] for r in cur.fetchall()) or "unknown"
|
|
|
|
# Date.
|
|
date_str = _time.strftime("%Y%m%d", _time.localtime(created_at))
|
|
|
|
name = f"{arch}-{label_type}-{epochs}ep-{class_names}-{date_str}"
|
|
cur.execute("UPDATE model_versions SET name = ? WHERE id = ?", (name, model_id))
|
|
|
|
|
|
def _backfill_dataset_rules(cur) -> None:
|
|
"""Datasets merged before REQ-132 have no snapshot. Give them the project's
|
|
current rules — that is what they were cut under, unless the rules changed
|
|
in between, and an empty snapshot would claim "no rules" for all of them.
|
|
"""
|
|
import json
|
|
|
|
cur.execute("SELECT id, project_id FROM datasets WHERE rules_json IS NULL")
|
|
stale = cur.fetchall()
|
|
if not stale:
|
|
return
|
|
by_project: dict = {}
|
|
for dataset_id, project_id in stale:
|
|
if project_id not in by_project:
|
|
cur.execute(
|
|
"""SELECT name, predicate, action, target_class FROM triage_rules
|
|
WHERE project_id = ? AND stage = 'dataprep' ORDER BY position""",
|
|
(project_id,),
|
|
)
|
|
by_project[project_id] = json.dumps([
|
|
{"name": row[0], "predicate": json.loads(row[1]),
|
|
"action": row[2], "target_class": row[3]}
|
|
for row in cur.fetchall()
|
|
])
|
|
cur.execute("UPDATE datasets SET rules_json = ? WHERE id = ?",
|
|
(by_project[project_id], dataset_id))
|
|
|
|
|
|
def healthy() -> bool:
|
|
try:
|
|
with cursor() as cur:
|
|
cur.execute("SELECT 1")
|
|
return True
|
|
except sqlite3.Error:
|
|
return False
|