mirror of
https://github.com/MrSleeps/jellyfin-scene-search.git
synced 2026-11-15 22:41:28 +00:00
Sleeps
132 lines
3.8 KiB
Python
132 lines
3.8 KiB
Python
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
from typing import Iterator
|
|
|
|
import numpy as np
|
|
|
|
SCHEMA = """
|
|
PRAGMA foreign_keys = ON;
|
|
PRAGMA journal_mode = WAL;
|
|
|
|
CREATE TABLE IF NOT EXISTS media (
|
|
id TEXT PRIMARY KEY,
|
|
source_path TEXT NOT NULL,
|
|
title TEXT NOT NULL,
|
|
content_sha256 TEXT NOT NULL,
|
|
size_bytes INTEGER NOT NULL,
|
|
mtime_ns INTEGER NOT NULL,
|
|
interval_ms INTEGER NOT NULL,
|
|
model_id TEXT NOT NULL,
|
|
embedding_dim INTEGER NOT NULL,
|
|
source_kind TEXT NOT NULL DEFAULT 'file',
|
|
jellyfin_item_id TEXT,
|
|
jellyfin_media_source_id TEXT,
|
|
jellyfin_library_id TEXT,
|
|
jellyfin_item_type TEXT,
|
|
jellyfin_tags TEXT,
|
|
jellyfin_etag TEXT,
|
|
trickplay_width INTEGER,
|
|
trickplay_tile_width INTEGER,
|
|
trickplay_tile_height INTEGER,
|
|
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS frames (
|
|
id TEXT PRIMARY KEY,
|
|
media_id TEXT NOT NULL REFERENCES media(id) ON DELETE CASCADE,
|
|
frame_number INTEGER NOT NULL,
|
|
timestamp_ms INTEGER NOT NULL,
|
|
thumb_path TEXT NOT NULL,
|
|
embedding_f16 BLOB NOT NULL,
|
|
caption TEXT,
|
|
caption_model TEXT,
|
|
captioned_at TEXT,
|
|
caption_route TEXT,
|
|
caption_route_reason TEXT,
|
|
route_score REAL,
|
|
route_model TEXT,
|
|
route_policy TEXT,
|
|
routed_at TEXT,
|
|
UNIQUE(media_id, frame_number)
|
|
);
|
|
|
|
CREATE INDEX IF NOT EXISTS idx_frames_media_timestamp
|
|
ON frames(media_id, timestamp_ms);
|
|
"""
|
|
|
|
|
|
@contextmanager
|
|
def connect(path: Path) -> Iterator[sqlite3.Connection]:
|
|
con = sqlite3.connect(path)
|
|
con.row_factory = sqlite3.Row
|
|
con.execute("PRAGMA foreign_keys = ON")
|
|
try:
|
|
yield con
|
|
con.commit()
|
|
finally:
|
|
con.close()
|
|
|
|
|
|
def _ensure_media_jellyfin_columns(con: sqlite3.Connection) -> None:
|
|
columns = {row["name"] for row in con.execute("PRAGMA table_info(media)")}
|
|
additions = {
|
|
"source_kind": "TEXT NOT NULL DEFAULT 'file'",
|
|
"jellyfin_item_id": "TEXT",
|
|
"jellyfin_media_source_id": "TEXT",
|
|
"jellyfin_library_id": "TEXT",
|
|
"jellyfin_item_type": "TEXT",
|
|
"jellyfin_tags": "TEXT",
|
|
"jellyfin_etag": "TEXT",
|
|
"trickplay_width": "INTEGER",
|
|
"trickplay_tile_width": "INTEGER",
|
|
"trickplay_tile_height": "INTEGER",
|
|
}
|
|
for name, sql_type in additions.items():
|
|
if name not in columns:
|
|
con.execute(f"ALTER TABLE media ADD COLUMN {name} {sql_type}")
|
|
con.execute(
|
|
"""
|
|
CREATE UNIQUE INDEX IF NOT EXISTS idx_media_jellyfin_item
|
|
ON media(jellyfin_item_id)
|
|
WHERE jellyfin_item_id IS NOT NULL
|
|
"""
|
|
)
|
|
|
|
|
|
def _ensure_frame_caption_columns(con: sqlite3.Connection) -> None:
|
|
columns = {row["name"] for row in con.execute("PRAGMA table_info(frames)")}
|
|
additions = {
|
|
"caption": "TEXT",
|
|
"caption_model": "TEXT",
|
|
"captioned_at": "TEXT",
|
|
"caption_route": "TEXT",
|
|
"caption_route_reason": "TEXT",
|
|
"route_score": "REAL",
|
|
"route_model": "TEXT",
|
|
"route_policy": "TEXT",
|
|
"routed_at": "TEXT",
|
|
}
|
|
for name, sql_type in additions.items():
|
|
if name not in columns:
|
|
con.execute(f"ALTER TABLE frames ADD COLUMN {name} {sql_type}")
|
|
|
|
|
|
def init_db(path: Path) -> None:
|
|
with connect(path) as con:
|
|
con.executescript(SCHEMA)
|
|
_ensure_media_jellyfin_columns(con)
|
|
_ensure_frame_caption_columns(con)
|
|
|
|
|
|
def pack_vector(vector: np.ndarray) -> bytes:
|
|
# Canonical store uses float16 to halve disk usage. Meilisearch receives float32.
|
|
return np.asarray(vector, dtype=np.float16).tobytes()
|
|
|
|
|
|
def unpack_vector(blob: bytes, dim: int) -> np.ndarray:
|
|
return np.frombuffer(blob, dtype=np.float16, count=dim).astype(np.float32)
|