179 lines
6.7 KiB
Python
179 lines
6.7 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 dataset_items (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
project_id INTEGER NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
|
|
frame_id INTEGER NOT NULL REFERENCES frames(id) ON DELETE CASCADE UNIQUE,
|
|
split TEXT NOT NULL CHECK (split IN ('train', 'val')),
|
|
image_rel TEXT NOT NULL,
|
|
label_rel TEXT NOT NULL,
|
|
added_at REAL NOT NULL
|
|
)
|
|
""",
|
|
"""
|
|
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')),
|
|
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 INDEX IF NOT EXISTS idx_frames_batch ON frames(batch_id, idx)",
|
|
"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)",
|
|
]
|
|
|
|
|
|
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")
|
|
|
|
|
|
def healthy() -> bool:
|
|
try:
|
|
with cursor() as cur:
|
|
cur.execute("SELECT 1")
|
|
return True
|
|
except sqlite3.Error:
|
|
return False
|