refactor: strip AI pipeline to binary photo/other classifier
Drops face recognition, OCR, object detection, and semantic embeddings. The sole remaining vision task is a CLIP-based binary classifier (photography vs other); photos in "other" get needs_review=true so screenshots, documents, memes and scans can be triaged from a new filter pill in the UI. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -41,12 +41,10 @@ import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Optional
|
||||
|
||||
from sqlalchemy import select, update, text
|
||||
from sqlalchemy import select, update
|
||||
|
||||
from app.database import AsyncSessionLocal
|
||||
from app.models.photos import Photo
|
||||
from app.models.embeddings import Embedding
|
||||
from app.config import settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -54,11 +52,6 @@ logger = logging.getLogger(__name__)
|
||||
# pHash Hamming distance threshold (6 out of 64 bits).
|
||||
DEFAULT_PHASH_THRESHOLD = 6
|
||||
|
||||
# CLIP cosine distance threshold. CLIP embeddings are L2-normalized,
|
||||
# so cosine distance = 1 - dot(a, b). A threshold of 0.08 catches
|
||||
# visually near-identical shots; 0.15 catches similar compositions.
|
||||
DEFAULT_CLIP_THRESHOLD = 0.10
|
||||
|
||||
|
||||
def _hex_to_int(h: str) -> int:
|
||||
"""Parse a 16-char hex pHash to a Python int. Returns -1 on bad input
|
||||
@@ -119,14 +112,12 @@ class _UnionFind:
|
||||
|
||||
async def regroup_duplicates(
|
||||
phash_threshold: int = DEFAULT_PHASH_THRESHOLD,
|
||||
clip_threshold: float = DEFAULT_CLIP_THRESHOLD,
|
||||
**_ignored,
|
||||
) -> dict:
|
||||
"""Full recompute of duplicate groups using pHash + CLIP similarity.
|
||||
"""Full recompute of duplicate groups using pHash similarity.
|
||||
|
||||
Idempotent — safe to call as often as you like. Returns a summary dict.
|
||||
"""
|
||||
embedder_model = settings.vision.embedder.name
|
||||
|
||||
async with AsyncSessionLocal() as session:
|
||||
# Pull all visible photos with a phash or embedding.
|
||||
rows = (
|
||||
@@ -162,15 +153,6 @@ async def regroup_duplicates(
|
||||
if _hamming(hi, hj) <= phash_threshold:
|
||||
uf.union_by_key(phash_ids[i], phash_ids[j])
|
||||
|
||||
# ── Phase 2: CLIP similarity via pgvector ──
|
||||
# For each photo with an embedding, find its nearest neighbors
|
||||
# within the cosine distance threshold using the HNSW index.
|
||||
clip_matches = await _clip_neighbor_scan(
|
||||
session, ids, embedder_model, clip_threshold
|
||||
)
|
||||
for photo_id, neighbor_id in clip_matches:
|
||||
uf.union_by_key(photo_id, neighbor_id)
|
||||
|
||||
# ── Write results ──
|
||||
await _clear_all_groups(session)
|
||||
|
||||
@@ -202,16 +184,9 @@ async def regroup_duplicates(
|
||||
async def incremental_regroup(
|
||||
since: Optional[datetime] = None,
|
||||
phash_threshold: int = DEFAULT_PHASH_THRESHOLD,
|
||||
clip_threshold: float = DEFAULT_CLIP_THRESHOLD,
|
||||
**_ignored,
|
||||
) -> dict:
|
||||
"""Incremental duplicate detection for newly added photos.
|
||||
|
||||
Only photos added after `since` are compared against the full library.
|
||||
Much faster than a full regroup for post-scan updates:
|
||||
O(new × log N) via HNSW instead of O(N²).
|
||||
"""
|
||||
embedder_model = settings.vision.embedder.name
|
||||
|
||||
"""Incremental duplicate detection for newly added photos using pHash."""
|
||||
async with AsyncSessionLocal() as session:
|
||||
# If no watermark, fall back to full regroup.
|
||||
if since is None:
|
||||
@@ -282,13 +257,6 @@ async def incremental_regroup(
|
||||
if _hamming(nh, eh) <= phash_threshold:
|
||||
uf.union_by_key(new_id, existing_id)
|
||||
|
||||
# ── Phase 2: CLIP — vector similarity for new photos only ──
|
||||
clip_matches = await _clip_neighbor_scan(
|
||||
session, new_ids, embedder_model, clip_threshold
|
||||
)
|
||||
for photo_id, neighbor_id in clip_matches:
|
||||
uf.union_by_key(photo_id, neighbor_id)
|
||||
|
||||
# ── Write results ──
|
||||
# Only update groups that contain at least one new photo.
|
||||
# Clear all groups first, then rewrite.
|
||||
@@ -323,73 +291,6 @@ async def incremental_regroup(
|
||||
}
|
||||
|
||||
|
||||
async def _clip_neighbor_scan(
|
||||
session,
|
||||
photo_ids: list[str],
|
||||
embedder_model: str,
|
||||
threshold: float,
|
||||
) -> list[tuple[str, str]]:
|
||||
"""For each photo in `photo_ids` that has a CLIP embedding, find
|
||||
neighbors within cosine distance `threshold` using pgvector HNSW.
|
||||
|
||||
Returns a list of (photo_id, neighbor_id) pairs.
|
||||
"""
|
||||
matches: list[tuple[str, str]] = []
|
||||
|
||||
if not photo_ids:
|
||||
return matches
|
||||
|
||||
# Batch: get all embeddings for the target photos.
|
||||
target_embeddings = (
|
||||
await session.execute(
|
||||
select(Embedding.photo_id, Embedding.vector)
|
||||
.where(Embedding.photo_id.in_(photo_ids))
|
||||
.where(Embedding.model == embedder_model)
|
||||
)
|
||||
).all()
|
||||
|
||||
if not target_embeddings:
|
||||
return matches
|
||||
|
||||
# For each target, query nearest neighbors via pgvector.
|
||||
# We use raw SQL for the <=> cosine distance operator.
|
||||
for photo_id, vector in target_embeddings:
|
||||
# pgvector cosine distance: <=> operator
|
||||
# Find top 20 nearest neighbors within threshold.
|
||||
# Serialize the vector as "[a,b,c,...]" — pgvector's text
|
||||
# format uses commas; numpy's default str() joins with spaces
|
||||
# which Postgres rejects with "invalid input syntax for vector".
|
||||
if hasattr(vector, 'tolist'):
|
||||
vec_seq = vector.tolist()
|
||||
else:
|
||||
vec_seq = list(vector)
|
||||
vec_text = '[' + ','.join(f'{float(x):.8f}' for x in vec_seq) + ']'
|
||||
result = await session.execute(
|
||||
text("""
|
||||
SELECT e.photo_id, (e.vector <=> :vec) AS distance
|
||||
FROM embeddings e
|
||||
JOIN photos p ON p.id = e.photo_id
|
||||
WHERE e.model = :model
|
||||
AND e.photo_id != :pid
|
||||
AND p.is_trashed = false
|
||||
AND p.is_hidden = false
|
||||
AND (e.vector <=> :vec) < :threshold
|
||||
ORDER BY e.vector <=> :vec
|
||||
LIMIT 20
|
||||
"""),
|
||||
{
|
||||
'vec': vec_text,
|
||||
'pid': photo_id,
|
||||
'model': embedder_model,
|
||||
'threshold': threshold,
|
||||
}
|
||||
)
|
||||
for row in result.all():
|
||||
matches.append((photo_id, row[0]))
|
||||
|
||||
return matches
|
||||
|
||||
|
||||
async def _clear_all_groups(session) -> None:
|
||||
"""Reset duplicate_group_id / is_duplicate on every photo."""
|
||||
await session.execute(
|
||||
|
||||
@@ -1,23 +1,9 @@
|
||||
"""
|
||||
Runtime feature flags for expensive pipeline stages.
|
||||
Runtime feature flags for the vision pipeline.
|
||||
|
||||
The YAML config (``mulita.yml``) ships reasonable defaults. Admins can
|
||||
toggle these at runtime from the Settings → AI Features tab without
|
||||
rebuilding the image or editing the bind-mounted YAML; the overrides
|
||||
live in Redis so both the FastAPI backend and the Celery workers see
|
||||
the same value within ~1s of the write.
|
||||
|
||||
The key namespace is:
|
||||
|
||||
mulita:flags:<name> → "true" | "false"
|
||||
|
||||
An unset key means "fall back to the YAML default" — so an admin who
|
||||
has never touched the tab sees exactly the config-file behaviour.
|
||||
|
||||
Only bool flags live here. Thresholds, batch sizes, model names etc.
|
||||
stay in the YAML file because flipping them safely requires restarting
|
||||
the vision workers (model reload, ONNX session re-init); that's not
|
||||
something a single admin click should do.
|
||||
Only one flag now — the master vision switch. Runtime overrides live in
|
||||
Redis under ``mulita:flags:<name>``; an unset key falls back to the
|
||||
YAML default.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -31,30 +17,16 @@ from app.config import settings
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Feature identifiers. The public name is what the admin UI sends; the
|
||||
# ``yaml_default`` getter returns the value the YAML would have set.
|
||||
# Keep these in sync with the VisionSettings fields in ``config.py``.
|
||||
FLAG_VISION_ENABLED = 'vision.enabled'
|
||||
FLAG_OCR_ENABLED = 'vision.ocr.enabled'
|
||||
FLAG_DETECTOR_ENABLED = 'vision.detector.enabled'
|
||||
FLAG_FACES_ENABLED = 'vision.faces.enabled'
|
||||
FLAG_CLASSIFIER_ENABLED = 'vision.classifier.enabled'
|
||||
|
||||
ALL_FLAGS = (
|
||||
FLAG_VISION_ENABLED,
|
||||
FLAG_OCR_ENABLED,
|
||||
FLAG_DETECTOR_ENABLED,
|
||||
FLAG_FACES_ENABLED,
|
||||
FLAG_CLASSIFIER_ENABLED,
|
||||
)
|
||||
ALL_FLAGS = (FLAG_VISION_ENABLED,)
|
||||
|
||||
_VISION_QUEUE = 'vision'
|
||||
|
||||
_REDIS: Optional[redis.Redis] = None
|
||||
|
||||
|
||||
def _redis() -> Optional[redis.Redis]:
|
||||
"""Lazy Redis client. Returns None if the broker is unreachable so
|
||||
callers can fall back to YAML defaults instead of crashing."""
|
||||
global _REDIS
|
||||
if _REDIS is None:
|
||||
try:
|
||||
@@ -69,19 +41,8 @@ def _redis() -> Optional[redis.Redis]:
|
||||
|
||||
|
||||
def _yaml_default(name: str) -> bool:
|
||||
"""Return the YAML-configured default for a flag. Used when Redis
|
||||
has no value set (fresh install or admin never touched the tab)."""
|
||||
v = settings.vision
|
||||
if name == FLAG_VISION_ENABLED:
|
||||
return bool(v.enabled)
|
||||
if name == FLAG_OCR_ENABLED:
|
||||
return bool(v.ocr.enabled)
|
||||
if name == FLAG_DETECTOR_ENABLED:
|
||||
return bool(v.detector.enabled)
|
||||
if name == FLAG_FACES_ENABLED:
|
||||
return bool(v.faces.enabled)
|
||||
if name == FLAG_CLASSIFIER_ENABLED:
|
||||
return bool(v.classifier.enabled)
|
||||
return bool(settings.vision.enabled)
|
||||
raise ValueError(f"Unknown feature flag: {name!r}")
|
||||
|
||||
|
||||
@@ -90,16 +51,6 @@ def _redis_key(name: str) -> str:
|
||||
|
||||
|
||||
def is_enabled(name: str) -> bool:
|
||||
"""Return True if feature ``name`` is currently enabled.
|
||||
|
||||
Order of precedence:
|
||||
1. Redis override (set by PATCH /admin/feature-flags)
|
||||
2. YAML default
|
||||
|
||||
Reads are cheap (~ms) and we intentionally do NOT add a local
|
||||
process cache — the whole point of runtime flags is that a toggle
|
||||
takes effect on the next task without a worker restart.
|
||||
"""
|
||||
r = _redis()
|
||||
if r is not None:
|
||||
try:
|
||||
@@ -112,9 +63,6 @@ def is_enabled(name: str) -> bool:
|
||||
|
||||
|
||||
def set_flag(name: str, value: bool) -> None:
|
||||
"""Persist a flag override to Redis. No-op if Redis is unreachable
|
||||
(we don't silently pretend to have written; raise so the admin
|
||||
request returns a 500 instead of misleading success)."""
|
||||
if name not in ALL_FLAGS:
|
||||
raise ValueError(f"Unknown feature flag: {name!r}")
|
||||
r = _redis()
|
||||
@@ -125,9 +73,6 @@ def set_flag(name: str, value: bool) -> None:
|
||||
|
||||
|
||||
def reset_flag(name: str) -> None:
|
||||
"""Delete the Redis override so the flag falls back to its YAML
|
||||
default. Useful if an admin wants a clean slate without guessing
|
||||
what the config defaults are."""
|
||||
if name not in ALL_FLAGS:
|
||||
raise ValueError(f"Unknown feature flag: {name!r}")
|
||||
r = _redis()
|
||||
@@ -137,150 +82,41 @@ def reset_flag(name: str) -> None:
|
||||
_apply_worker_side_effects(name)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Worker-level side effects: when the admin flips a flag we don't just want
|
||||
# gating at task-start (which still executes the message, it just returns
|
||||
# 'skipped'). We also want queued work gone and the vision worker genuinely
|
||||
# idle when the master switch is off.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_VISION_QUEUE = 'vision'
|
||||
|
||||
# Flag → celery task name(s) whose queued messages should be dropped when
|
||||
# the flag goes off. Keeps the queue from replaying yesterday's work the
|
||||
# moment someone re-enables the stage.
|
||||
_TASKS_BY_FLAG: dict[str, tuple[str, ...]] = {
|
||||
FLAG_VISION_ENABLED: (
|
||||
'embed_photo', 'ocr_photo', 'detect_objects', 'extract_faces',
|
||||
'classify_content', 'vision_fanout', 'recluster_faces',
|
||||
),
|
||||
FLAG_OCR_ENABLED: ('ocr_photo',),
|
||||
FLAG_DETECTOR_ENABLED: ('detect_objects',),
|
||||
FLAG_FACES_ENABLED: ('extract_faces', 'recluster_faces'),
|
||||
FLAG_CLASSIFIER_ENABLED: ('classify_content',),
|
||||
}
|
||||
|
||||
|
||||
def _apply_worker_side_effects(name: str) -> None:
|
||||
"""Bring the live workers in line with the new flag value.
|
||||
|
||||
For the master ``vision.enabled`` flag we go beyond task gating and
|
||||
actually stop consumption from the ``vision`` queue — flipping it
|
||||
off puts the vision worker to sleep (no CPU, no model memory
|
||||
churn) until it's flipped back on. For per-feature flags, the
|
||||
running tasks already skip via ``is_enabled``; we just purge any
|
||||
messages already sitting in the queue so the admin doesn't pay for
|
||||
a backlog on re-enable.
|
||||
|
||||
All operations are best-effort — if control messaging or a Redis
|
||||
op fails, we log and return; the flag state itself is already
|
||||
persisted so the gating path continues to work.
|
||||
"""
|
||||
"""Attach or detach the vision consumer and purge queued work when
|
||||
the master flag flips. Best-effort — state is already persisted."""
|
||||
if name != FLAG_VISION_ENABLED:
|
||||
return
|
||||
try:
|
||||
# Lazy import: avoids a circular dependency between the services
|
||||
# module (imported from tasks.vision) and the celery app config.
|
||||
from app.tasks.celery import celery_app
|
||||
except Exception as e:
|
||||
logger.warning(f"feature_flags: celery app unavailable for side effects ({e})")
|
||||
return
|
||||
|
||||
try:
|
||||
if name == FLAG_VISION_ENABLED:
|
||||
if is_enabled(FLAG_VISION_ENABLED):
|
||||
# Re-attach the vision consumer so workers pick up tasks
|
||||
# again. broadcast=True ensures every running worker
|
||||
# receives the command.
|
||||
celery_app.control.add_consumer(_VISION_QUEUE, reply=False)
|
||||
logger.info("feature_flags: vision re-enabled; consumer added")
|
||||
else:
|
||||
celery_app.control.cancel_consumer(_VISION_QUEUE, reply=False)
|
||||
_purge_queue(_VISION_QUEUE)
|
||||
logger.info(
|
||||
"feature_flags: vision disabled; consumer cancelled "
|
||||
"and queue purged"
|
||||
)
|
||||
return
|
||||
|
||||
# Per-feature flag going off → drop pending tasks of its types.
|
||||
if not is_enabled(name):
|
||||
targets = _TASKS_BY_FLAG.get(name, ())
|
||||
if targets:
|
||||
removed = _purge_queue_by_task_names(_VISION_QUEUE, targets)
|
||||
logger.info(
|
||||
f"feature_flags: {name} disabled; removed {removed} "
|
||||
f"pending messages from {_VISION_QUEUE}"
|
||||
)
|
||||
if is_enabled(FLAG_VISION_ENABLED):
|
||||
celery_app.control.add_consumer(_VISION_QUEUE, reply=False)
|
||||
logger.info("feature_flags: vision re-enabled; consumer added")
|
||||
else:
|
||||
celery_app.control.cancel_consumer(_VISION_QUEUE, reply=False)
|
||||
_purge_queue(_VISION_QUEUE)
|
||||
logger.info("feature_flags: vision disabled; consumer cancelled and queue purged")
|
||||
except Exception as e:
|
||||
logger.warning(f"feature_flags: worker side effects failed for {name}: {e}")
|
||||
logger.warning(f"feature_flags: worker side effects failed: {e}")
|
||||
|
||||
|
||||
def _purge_queue(queue: str) -> int:
|
||||
"""Drop every pending message from ``queue``. Returns the count
|
||||
deleted. Celery's control.purge() purges the default queue only,
|
||||
so we delete the Redis key directly (the broker's queue list)."""
|
||||
r = _redis()
|
||||
if r is None:
|
||||
return 0
|
||||
try:
|
||||
removed = r.delete(queue)
|
||||
return int(removed or 0)
|
||||
return int(r.delete(queue) or 0)
|
||||
except Exception as e:
|
||||
logger.warning(f"feature_flags: purge {queue} failed: {e}")
|
||||
return 0
|
||||
|
||||
|
||||
def _purge_queue_by_task_names(queue: str, task_names: tuple[str, ...]) -> int:
|
||||
"""Walk ``queue`` and drop any message whose Celery task name is in
|
||||
``task_names``. Other messages are preserved (pushed back in order)
|
||||
so we don't flush embed tasks when the admin disabled only OCR.
|
||||
|
||||
Celery stores each message as a JSON blob in a Redis list; the
|
||||
task name lives at ``headers.task``.
|
||||
"""
|
||||
import json
|
||||
r = _redis()
|
||||
if r is None:
|
||||
return 0
|
||||
try:
|
||||
# Snapshot the queue, then rebuild it without the filtered names.
|
||||
# Done inside a Redis transaction so a concurrent enqueue doesn't
|
||||
# race with us (worst case it gets re-delivered after we release,
|
||||
# which is the normal enqueue path anyway).
|
||||
pipe = r.pipeline()
|
||||
pipe.lrange(queue, 0, -1)
|
||||
pipe.delete(queue)
|
||||
raw_items, _ = pipe.execute()
|
||||
kept: list[bytes | str] = []
|
||||
removed = 0
|
||||
for raw in raw_items or []:
|
||||
try:
|
||||
# Messages can be bytes or str depending on decode_responses.
|
||||
payload = raw.decode() if isinstance(raw, bytes) else raw
|
||||
msg = json.loads(payload)
|
||||
task = (
|
||||
msg.get('headers', {}).get('task')
|
||||
or msg.get('task')
|
||||
)
|
||||
if task in task_names:
|
||||
removed += 1
|
||||
continue
|
||||
except Exception:
|
||||
# Unparseable message — keep it, better to leak than
|
||||
# to silently drop a message we can't identify.
|
||||
pass
|
||||
kept.append(raw)
|
||||
if kept:
|
||||
r.rpush(queue, *kept)
|
||||
return removed
|
||||
except Exception as e:
|
||||
logger.warning(f"feature_flags: selective purge failed on {queue}: {e}")
|
||||
return 0
|
||||
|
||||
|
||||
def snapshot() -> dict[str, dict[str, object]]:
|
||||
"""Return every flag's current effective value, YAML default, and
|
||||
whether it's overridden. Powers the admin UI tab.
|
||||
"""
|
||||
r = _redis()
|
||||
out: dict[str, dict[str, object]] = {}
|
||||
for name in ALL_FLAGS:
|
||||
|
||||
@@ -1,19 +1,13 @@
|
||||
"""
|
||||
Unified search service — hybrid FTS + semantic (RRF) search.
|
||||
|
||||
Phase 1 (PR4): semantic-only via pgvector cosine similarity.
|
||||
Phase 2 (PR5): adds FTS via tsvector, enables RRF fusion.
|
||||
FTS search over photos.search_vector with optional tag/date filters.
|
||||
"""
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
from sqlalchemy import select, text, func
|
||||
from sqlalchemy import select, text
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.models import Photo
|
||||
from app.models.embeddings import Embedding
|
||||
from app.config import settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -27,135 +21,57 @@ async def hybrid_search(
|
||||
limit: int = 50,
|
||||
offset: int = 0,
|
||||
) -> list[dict]:
|
||||
"""Run hybrid search (FTS + semantic) with RRF fusion.
|
||||
|
||||
Currently semantic-only; FTS leg added in PR5.
|
||||
"""
|
||||
model_name = settings.vision.embedder.name
|
||||
results = {}
|
||||
|
||||
# ── Semantic search (CLIP text → pgvector cosine) ─────────────────
|
||||
"""Full-text search using photos.search_vector. No embeddings, no OCR."""
|
||||
if q:
|
||||
try:
|
||||
from app.services.vision.registry import registry
|
||||
embedder = registry.get_embedder()
|
||||
query_vec = embedder.embed_text(q)
|
||||
|
||||
# pgvector cosine distance: <=> returns distance (lower = closer).
|
||||
# Join photos so we can filter out discarded / hidden rows
|
||||
# inside the same query — otherwise a hidden-folder photo
|
||||
# can take a top-N rank and starve the visible results.
|
||||
vec_str = "[" + ",".join(str(float(v)) for v in query_vec) + "]"
|
||||
stmt = text("""
|
||||
SELECT e.photo_id,
|
||||
(e.vector <=> :qvec::vector) AS distance
|
||||
FROM embeddings e
|
||||
JOIN photos p ON p.id = e.photo_id
|
||||
WHERE e.model = :model
|
||||
AND p.is_trashed = false
|
||||
AND p.is_hidden = false
|
||||
ORDER BY e.vector <=> :qvec::vector
|
||||
LIMIT 200
|
||||
""")
|
||||
rows = (await db.execute(stmt, {"qvec": vec_str, "model": model_name})).fetchall()
|
||||
|
||||
for rank, (photo_id, distance) in enumerate(rows):
|
||||
if photo_id not in results:
|
||||
results[photo_id] = {"semantic_rank": rank, "fts_rank": None}
|
||||
else:
|
||||
results[photo_id]["semantic_rank"] = rank
|
||||
|
||||
except Exception as e:
|
||||
logger.warning("Semantic search failed (models may not be loaded): %s", e)
|
||||
|
||||
# ── FTS search (photos.search_vector + ocr_text) ────────────────
|
||||
if q:
|
||||
try:
|
||||
# Same discarded/hidden filter as the semantic leg.
|
||||
# The OCR branch joins photos (through photo_id) so we can
|
||||
# filter there too; otherwise OCR hits in hidden folders
|
||||
# would leak into results.
|
||||
fts_stmt = text("""
|
||||
SELECT id, ts_rank(search_vector, plainto_tsquery('english', :q)) AS rank
|
||||
FROM photos
|
||||
WHERE search_vector @@ plainto_tsquery('english', :q)
|
||||
AND is_trashed = false
|
||||
AND is_hidden = false
|
||||
UNION
|
||||
SELECT o.photo_id AS id,
|
||||
MAX(o.confidence) AS rank
|
||||
FROM ocr_text o
|
||||
JOIN photos p ON p.id = o.photo_id
|
||||
WHERE to_tsvector('english', o.text) @@ plainto_tsquery('english', :q)
|
||||
AND p.is_trashed = false
|
||||
AND p.is_hidden = false
|
||||
GROUP BY o.photo_id
|
||||
ORDER BY rank DESC
|
||||
LIMIT 200
|
||||
LIMIT 500
|
||||
""")
|
||||
fts_rows = (await db.execute(fts_stmt, {"q": q})).fetchall()
|
||||
for rank, (photo_id, score) in enumerate(fts_rows):
|
||||
if photo_id not in results:
|
||||
results[photo_id] = {"semantic_rank": None, "fts_rank": rank}
|
||||
else:
|
||||
results[photo_id]["fts_rank"] = rank
|
||||
rows = (await db.execute(fts_stmt, {"q": q})).fetchall()
|
||||
except Exception as e:
|
||||
logger.warning("FTS search failed: %s", e)
|
||||
rows = []
|
||||
|
||||
# ── RRF fusion ────────────────────────────────────────────────────
|
||||
k = 60
|
||||
scored = []
|
||||
for photo_id, ranks in results.items():
|
||||
score = 0.0
|
||||
if ranks["semantic_rank"] is not None:
|
||||
score += 1.0 / (k + ranks["semantic_rank"])
|
||||
if ranks.get("fts_rank") is not None:
|
||||
score += 1.0 / (k + ranks["fts_rank"])
|
||||
scored.append((photo_id, score))
|
||||
scored = [(pid, float(rank)) for pid, rank in rows]
|
||||
|
||||
scored.sort(key=lambda x: -x[1])
|
||||
|
||||
# If no text query, fall back to recent photos. Always filter out
|
||||
# discarded + hidden here — this path backs the "Tags" and "People"
|
||||
# browse views, which should honor the folder hide flag.
|
||||
if not q:
|
||||
if tag_ids:
|
||||
from app.models.tags import photo_tags
|
||||
# Subquery to get distinct photo_ids matching the tag filter
|
||||
sub = select(photo_tags.c.photo_id).where(
|
||||
photo_tags.c.tag_id.in_(tag_ids)
|
||||
).distinct().subquery()
|
||||
stmt = select(Photo.id).join(sub, Photo.id == sub.c.photo_id)
|
||||
else:
|
||||
stmt = select(Photo.id)
|
||||
stmt = stmt.where(
|
||||
Photo.is_discarded.is_(False),
|
||||
Photo.is_hidden.is_(False),
|
||||
)
|
||||
stmt = stmt.order_by(Photo.added_at.desc())
|
||||
if date_from:
|
||||
stmt = stmt.where(Photo.taken_at >= date_from)
|
||||
if date_to:
|
||||
stmt = stmt.where(Photo.taken_at <= date_to)
|
||||
stmt = stmt.offset(offset).limit(limit)
|
||||
rows = (await db.execute(stmt)).fetchall()
|
||||
return [{"photo_id": row[0], "score": 0.0} for row in rows]
|
||||
photo_ids = [pid for pid, _ in scored]
|
||||
if not photo_ids:
|
||||
return []
|
||||
stmt = select(photo_tags.c.photo_id).where(
|
||||
photo_tags.c.photo_id.in_(photo_ids),
|
||||
photo_tags.c.tag_id.in_(tag_ids),
|
||||
).distinct()
|
||||
valid = {row[0] for row in (await db.execute(stmt)).fetchall()}
|
||||
scored = [(pid, s) for pid, s in scored if pid in valid]
|
||||
|
||||
# Apply filters to scored results
|
||||
photo_ids = [pid for pid, _ in scored]
|
||||
if not photo_ids:
|
||||
return []
|
||||
page = scored[offset : offset + limit]
|
||||
return [{"photo_id": pid, "score": s} for pid, s in page]
|
||||
|
||||
# Filter by tags if requested
|
||||
# No text query — recent photos with tag/date filters.
|
||||
if tag_ids:
|
||||
from app.models.tags import photo_tags
|
||||
stmt = select(photo_tags.c.photo_id).where(
|
||||
photo_tags.c.photo_id.in_(photo_ids),
|
||||
photo_tags.c.tag_id.in_(tag_ids),
|
||||
).distinct()
|
||||
valid_ids = {row[0] for row in (await db.execute(stmt)).fetchall()}
|
||||
scored = [(pid, s) for pid, s in scored if pid in valid_ids]
|
||||
|
||||
# Paginate
|
||||
page = scored[offset : offset + limit]
|
||||
return [{"photo_id": pid, "score": score} for pid, score in page]
|
||||
sub = select(photo_tags.c.photo_id).where(
|
||||
photo_tags.c.tag_id.in_(tag_ids)
|
||||
).distinct().subquery()
|
||||
stmt = select(Photo.id).join(sub, Photo.id == sub.c.photo_id)
|
||||
else:
|
||||
stmt = select(Photo.id)
|
||||
stmt = stmt.where(
|
||||
Photo.is_discarded.is_(False),
|
||||
Photo.is_hidden.is_(False),
|
||||
)
|
||||
if date_from:
|
||||
stmt = stmt.where(Photo.taken_at >= date_from)
|
||||
if date_to:
|
||||
stmt = stmt.where(Photo.taken_at <= date_to)
|
||||
stmt = stmt.order_by(Photo.added_at.desc()).offset(offset).limit(limit)
|
||||
rows = (await db.execute(stmt)).fetchall()
|
||||
return [{"photo_id": row[0], "score": 0.0} for row in rows]
|
||||
|
||||
@@ -1,105 +1,25 @@
|
||||
"""
|
||||
Abstract base classes for vision backends.
|
||||
Abstract base classes for the vision backend.
|
||||
|
||||
Each ABC defines the contract a backend must satisfy. The default
|
||||
implementation is ONNXBackend (onnx_backend.py). A ROCm backend can be
|
||||
added later by subclassing these ABCs and registering via
|
||||
settings.vision.backend.
|
||||
The pipeline is now a single binary classifier: photography vs other.
|
||||
Feature extraction is an internal detail of the classifier and is not
|
||||
exposed as a separate service.
|
||||
"""
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
@dataclass
|
||||
class DetectionBox:
|
||||
"""A single object detection result."""
|
||||
label: str
|
||||
confidence: float
|
||||
bbox: list[float] # [x1, y1, x2, y2] normalized 0-1
|
||||
|
||||
|
||||
@dataclass
|
||||
class OCRResult:
|
||||
"""A single OCR text region."""
|
||||
text: str
|
||||
confidence: float
|
||||
bbox: list[float] # [x1, y1, x2, y2] normalized 0-1
|
||||
language: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class FaceDetection:
|
||||
"""A detected face with its recognition embedding."""
|
||||
bbox: list[float] # [x1, y1, x2, y2] normalized 0-1
|
||||
embedding: np.ndarray # float32 vector (128-d for SFace)
|
||||
quality: float
|
||||
|
||||
|
||||
@dataclass
|
||||
class ClassificationResult:
|
||||
"""A content-type classification."""
|
||||
label: str
|
||||
confidence: float
|
||||
|
||||
|
||||
class Embedder(ABC):
|
||||
"""Generates image and text embeddings (e.g. OpenCLIP ViT-B/32)."""
|
||||
|
||||
@abstractmethod
|
||||
def embed_image(self, image: np.ndarray) -> np.ndarray:
|
||||
"""Return a normalized float32 embedding vector for an RGB image."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def embed_text(self, text: str) -> np.ndarray:
|
||||
"""Return a normalized float32 embedding vector for a text query."""
|
||||
...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def dim(self) -> int:
|
||||
"""Dimensionality of the output embedding."""
|
||||
...
|
||||
|
||||
|
||||
class OCREngine(ABC):
|
||||
"""Extracts text from images (e.g. rapidocr-onnxruntime)."""
|
||||
|
||||
@abstractmethod
|
||||
def run(self, image: np.ndarray) -> list[OCRResult]:
|
||||
"""Return OCR results for an RGB image."""
|
||||
...
|
||||
|
||||
|
||||
class ObjectDetector(ABC):
|
||||
"""Detects objects in images (e.g. YOLOv8n)."""
|
||||
|
||||
@abstractmethod
|
||||
def detect(self, image: np.ndarray) -> list[DetectionBox]:
|
||||
"""Return detections for an RGB image."""
|
||||
...
|
||||
|
||||
|
||||
class ContentClassifier(ABC):
|
||||
"""Classifies images into content types (screenshot, document, etc.)."""
|
||||
"""Classifies an image into 'photography' or 'other'."""
|
||||
|
||||
@abstractmethod
|
||||
def classify(self, image: np.ndarray) -> list[ClassificationResult]:
|
||||
"""Return content type classifications for an RGB image."""
|
||||
...
|
||||
|
||||
|
||||
class FaceProcessor(ABC):
|
||||
"""Detects faces and extracts recognition embeddings (e.g. YuNet + SFace)."""
|
||||
|
||||
@abstractmethod
|
||||
def process(self, image: np.ndarray) -> list[FaceDetection]:
|
||||
"""Return face detections with embeddings for an RGB image."""
|
||||
...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def embedding_dim(self) -> int:
|
||||
"""Dimensionality of face embedding vectors."""
|
||||
def classify(self, image: np.ndarray) -> ClassificationResult:
|
||||
...
|
||||
|
||||
@@ -1,124 +1,46 @@
|
||||
"""
|
||||
Download vision model weights on first worker boot.
|
||||
|
||||
Run as: python -m app.services.vision.bootstrap_models
|
||||
|
||||
Or called from the vision worker entrypoint before Celery starts.
|
||||
Downloads are idempotent — existing files are skipped.
|
||||
|
||||
For models that require export (OpenCLIP, YOLOv8n), see export_models.py.
|
||||
Those must be exported once on any machine with pip, then placed in
|
||||
the models volume before the worker starts.
|
||||
Ensure the OpenCLIP ViT-B/32 visual encoder is present on worker boot.
|
||||
Exported via export_models.py if missing.
|
||||
"""
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from urllib.request import urlretrieve
|
||||
|
||||
from app.config import settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# (relative_path, url, description)
|
||||
# Models with url=None must be pre-exported via export_models.py.
|
||||
# InsightFace (RetinaFace + ArcFace) auto-downloads via the insightface
|
||||
# package on first use — no manual download entries needed.
|
||||
DOWNLOADS = []
|
||||
|
||||
# Models that need manual export via export_models.py
|
||||
EXPORTS = [
|
||||
REQUIRED = [
|
||||
("embed/visual.onnx", "OpenCLIP ViT-B/32 visual encoder"),
|
||||
("embed/textual.onnx", "OpenCLIP ViT-B/32 textual encoder"),
|
||||
("detect/yolov8n.onnx", "YOLOv8n object detector"),
|
||||
]
|
||||
|
||||
|
||||
def bootstrap(models_dir: str | None = None):
|
||||
"""Ensure all model files are present. Download what we can, warn about
|
||||
files that need manual export."""
|
||||
base = Path(models_dir or settings.vision.models_dir)
|
||||
base.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Download auto-downloadable models
|
||||
for rel_path, url, desc in DOWNLOADS:
|
||||
dest = base / rel_path
|
||||
dest.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if dest.exists():
|
||||
logger.debug("Already exists: %s (%s)", dest, desc)
|
||||
continue
|
||||
|
||||
logger.info("Downloading %s → %s", desc, dest)
|
||||
try:
|
||||
urlretrieve(url, str(dest))
|
||||
size_kb = dest.stat().st_size / 1024
|
||||
logger.info("Downloaded %s (%.0f KB)", desc, size_kb)
|
||||
except Exception as e:
|
||||
logger.error("Failed to download %s: %s", desc, e)
|
||||
if dest.exists():
|
||||
dest.unlink()
|
||||
|
||||
# Check for manually-exported models
|
||||
missing = []
|
||||
for rel_path, desc in EXPORTS:
|
||||
dest = base / rel_path
|
||||
if not dest.exists():
|
||||
missing.append((rel_path, desc))
|
||||
|
||||
missing = [(rel, desc) for rel, desc in REQUIRED if not (base / rel).exists()]
|
||||
if missing:
|
||||
logger.warning(
|
||||
"Missing %d model file(s); attempting automatic export:",
|
||||
len(missing),
|
||||
)
|
||||
for rel_path, desc in missing:
|
||||
logger.warning(" %s — %s", base / rel_path, desc)
|
||||
logger.warning("Missing %d model file(s); attempting automatic export", len(missing))
|
||||
try:
|
||||
from app.services.vision import export_models
|
||||
export_models.export_openclip_visual(base)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Export failed: %s. Run `python -m app.services.vision.export_models "
|
||||
"--models-dir %s` manually to retry.",
|
||||
e, base,
|
||||
)
|
||||
|
||||
from app.services.vision import export_models
|
||||
|
||||
missing_paths = {rel for rel, _ in missing}
|
||||
|
||||
# Only run the export functions whose outputs are actually missing.
|
||||
# Each function is mapped to the file(s) it produces.
|
||||
export_map = [
|
||||
(export_models.export_openclip, ["embed/visual.onnx", "embed/textual.onnx"]),
|
||||
(export_models.export_siglip2, ["embed_siglip2/visual.onnx", "embed_siglip2/textual.onnx"]),
|
||||
(export_models.export_yolov8n, ["detect/yolov8n.onnx"]),
|
||||
]
|
||||
for export_fn, outputs in export_map:
|
||||
if not any(o in missing_paths for o in outputs):
|
||||
continue
|
||||
try:
|
||||
export_fn(base)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Export step %s failed: %s. "
|
||||
"Run `python -m app.services.vision.export_models "
|
||||
"--models-dir %s` manually to retry.",
|
||||
export_fn.__name__,
|
||||
e,
|
||||
base,
|
||||
)
|
||||
|
||||
# Re-check what's still missing after the export pass.
|
||||
still_missing = [
|
||||
(rel_path, desc)
|
||||
for rel_path, desc in EXPORTS
|
||||
if not (base / rel_path).exists()
|
||||
]
|
||||
if still_missing:
|
||||
for rel_path, desc in still_missing:
|
||||
logger.error(" still missing: %s — %s", base / rel_path, desc)
|
||||
else:
|
||||
logger.info("All model files present in %s", base)
|
||||
still_missing = [(r, d) for r, d in REQUIRED if not (base / r).exists()]
|
||||
if still_missing:
|
||||
for rel, desc in still_missing:
|
||||
logger.error(" still missing: %s — %s", base / rel, desc)
|
||||
else:
|
||||
logger.info("All model files present in %s", base)
|
||||
|
||||
# Signal readiness via Redis so the scan pipeline knows the vision
|
||||
# worker can accept tasks.
|
||||
try:
|
||||
import redis as _redis
|
||||
r = _redis.from_url(settings.redis_url)
|
||||
r.set("mulita:vision:ready", "1")
|
||||
_redis.from_url(settings.redis_url).set("mulita:vision:ready", "1")
|
||||
logger.info("Set mulita:vision:ready in Redis")
|
||||
except Exception as e:
|
||||
logger.warning("Could not set vision readiness flag in Redis: %s", e)
|
||||
|
||||
@@ -1,117 +1,106 @@
|
||||
"""
|
||||
CLIP zero-shot content-type classifier.
|
||||
Binary content classifier: 'photography' vs 'other'.
|
||||
|
||||
Uses the native OpenCLIP PyTorch text encoder for high-quality text
|
||||
embeddings (the ONNX text encoder has degraded quality due to the
|
||||
eot_indices workaround). Image embeddings use the ONNX visual encoder
|
||||
which works well.
|
||||
Uses OpenCLIP ViT-B/32 image features (ONNX) and two pre-computed text
|
||||
prompt centroids. Text centroids are computed once with the native
|
||||
open_clip text encoder and cached to {models_dir}/classifier/vectors.npz
|
||||
so steady-state worker startup doesn't pay the PyTorch cost.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from app.config import VisionSettings
|
||||
from app.services.vision.base import ContentClassifier, ClassificationResult
|
||||
from app.services.vision.base import ClassificationResult, ContentClassifier
|
||||
from app.services.vision.embed import CLIPVisualEncoder
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
CATEGORY_PROMPTS = {
|
||||
"screenshot": [
|
||||
"a screenshot of a computer screen",
|
||||
"a screenshot of a phone screen",
|
||||
"a screen capture of a user interface",
|
||||
],
|
||||
"document": [
|
||||
"a scanned document",
|
||||
"a photo of a document with printed text",
|
||||
"a photo of a page of text on paper",
|
||||
],
|
||||
"receipt": [
|
||||
"a photo of a receipt",
|
||||
"a photo of a bill or invoice",
|
||||
],
|
||||
"meme": [
|
||||
"an internet meme with text overlay",
|
||||
"a funny image with caption text",
|
||||
],
|
||||
"artwork": [
|
||||
"a painting or drawing",
|
||||
"a sketch or illustration",
|
||||
"digital art or graphic design",
|
||||
],
|
||||
"photograph": [
|
||||
|
||||
PROMPTS = {
|
||||
"photography": [
|
||||
"a photograph taken with a camera",
|
||||
"a real photo of a real scene or person",
|
||||
"a candid photograph",
|
||||
"a portrait photograph",
|
||||
"a landscape photograph",
|
||||
],
|
||||
"other": [
|
||||
"a screenshot of a computer screen",
|
||||
"a screenshot of a phone screen",
|
||||
"a screen capture of a user interface",
|
||||
"a scanned document",
|
||||
"a photo of a document with printed text",
|
||||
"a photo of a receipt",
|
||||
"a photo of a bill or invoice",
|
||||
"an internet meme with text overlay",
|
||||
"a funny image with caption text",
|
||||
"a digital illustration or graphic design",
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _compute_text_centroids() -> dict[str, np.ndarray]:
|
||||
"""Compute the 'photography' and 'other' centroid vectors using the
|
||||
open_clip text encoder. Only called on the cache-miss path."""
|
||||
import open_clip
|
||||
import torch
|
||||
|
||||
logger.info("Computing CLIP text centroids for binary classifier")
|
||||
model, _, _ = open_clip.create_model_and_transforms(
|
||||
"ViT-B-32", pretrained="laion2b_s34b_b79k"
|
||||
)
|
||||
model.eval()
|
||||
tokenizer = open_clip.get_tokenizer("ViT-B-32")
|
||||
|
||||
centroids: dict[str, np.ndarray] = {}
|
||||
for label, prompts in PROMPTS.items():
|
||||
tokens = tokenizer(prompts)
|
||||
with torch.no_grad():
|
||||
feats = model.encode_text(tokens)
|
||||
feats = feats / feats.norm(dim=-1, keepdim=True)
|
||||
avg = feats.mean(dim=0)
|
||||
avg = avg / avg.norm()
|
||||
centroids[label] = avg.numpy().astype(np.float32)
|
||||
return centroids
|
||||
|
||||
|
||||
class CLIPContentClassifier(ContentClassifier):
|
||||
"""Zero-shot content classifier using CLIP text-image similarity.
|
||||
Uses native PyTorch for text encoding, ONNX for image encoding."""
|
||||
|
||||
def __init__(self, settings: VisionSettings):
|
||||
import open_clip
|
||||
|
||||
self._min_confidence = settings.classifier.min_confidence
|
||||
self._encoder = CLIPVisualEncoder(settings)
|
||||
|
||||
# Load native model for text encoding only.
|
||||
# Use whichever model family the embedder is configured for so
|
||||
# the classification text vectors live in the same space as the
|
||||
# image embeddings.
|
||||
embedder_name = settings.embedder.name
|
||||
if embedder_name.startswith("siglip2"):
|
||||
model_arch = "ViT-B-16-SigLIP-384"
|
||||
pretrained = "webli"
|
||||
cache_dir = Path(settings.models_dir) / "classifier"
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
cache_path = cache_dir / "vectors.npz"
|
||||
|
||||
if cache_path.exists():
|
||||
logger.info("Loading cached text centroids from %s", cache_path)
|
||||
data = np.load(cache_path)
|
||||
self._photo = data["photography"].astype(np.float32)
|
||||
self._other = data["other"].astype(np.float32)
|
||||
else:
|
||||
model_arch = "ViT-B-32"
|
||||
pretrained = "laion2b_s34b_b79k"
|
||||
centroids = _compute_text_centroids()
|
||||
self._photo = centroids["photography"]
|
||||
self._other = centroids["other"]
|
||||
np.savez(cache_path, photography=self._photo, other=self._other)
|
||||
logger.info("Cached text centroids to %s", cache_path)
|
||||
|
||||
logger.info("Loading %s text encoder for content classification", model_arch)
|
||||
model, _, _ = open_clip.create_model_and_transforms(
|
||||
model_arch, pretrained=pretrained
|
||||
)
|
||||
model.eval()
|
||||
self._model = model
|
||||
self._tokenizer = open_clip.get_tokenizer(model_arch)
|
||||
def classify(self, image: np.ndarray) -> ClassificationResult:
|
||||
vec = self._encoder.encode(image)
|
||||
s_photo = float(np.dot(vec, self._photo))
|
||||
s_other = float(np.dot(vec, self._other))
|
||||
|
||||
# Get the ONNX image embedder from the registry
|
||||
from app.services.vision.registry import registry
|
||||
self._embedder = registry.get_embedder()
|
||||
if s_photo >= s_other:
|
||||
label = "photography"
|
||||
margin = s_photo - s_other
|
||||
else:
|
||||
label = "other"
|
||||
margin = s_other - s_photo
|
||||
|
||||
# Pre-compute text embeddings for each category
|
||||
self._category_embeddings: dict[str, np.ndarray] = {}
|
||||
for category, prompts in CATEGORY_PROMPTS.items():
|
||||
tokens = self._tokenizer(prompts)
|
||||
with torch.no_grad():
|
||||
text_features = model.encode_text(tokens)
|
||||
text_features /= text_features.norm(dim=-1, keepdim=True)
|
||||
avg = text_features.mean(dim=0)
|
||||
avg /= avg.norm()
|
||||
self._category_embeddings[category] = avg.numpy().astype(np.float32)
|
||||
|
||||
logger.info("Content classifier ready with %d categories", len(self._category_embeddings))
|
||||
|
||||
def classify(self, image: np.ndarray) -> list[ClassificationResult]:
|
||||
img_vec = self._embedder.embed_image(image)
|
||||
|
||||
# Cosine similarity against each category
|
||||
scores = {}
|
||||
for category, cat_vec in self._category_embeddings.items():
|
||||
scores[category] = float(np.dot(img_vec, cat_vec))
|
||||
|
||||
# Sort by score descending
|
||||
ranked = sorted(scores.items(), key=lambda x: -x[1])
|
||||
best_cat, best_score = ranked[0]
|
||||
second_score = ranked[1][1]
|
||||
|
||||
margin = best_score - second_score
|
||||
# Normalize: 0.01 margin → ~0.5 confidence, 0.03+ → ~1.0
|
||||
# 0.01 margin → ~0.3 conf, 0.03+ → ~1.0
|
||||
confidence = min(1.0, margin * 30)
|
||||
|
||||
if confidence >= self._min_confidence:
|
||||
return [ClassificationResult(label=best_cat, confidence=confidence)]
|
||||
|
||||
return []
|
||||
return ClassificationResult(label=label, confidence=confidence)
|
||||
|
||||
@@ -1,42 +0,0 @@
|
||||
"""
|
||||
Face embedding clustering using DBSCAN with cosine distance.
|
||||
|
||||
Called by the periodic `recluster_faces` Celery task (PR7).
|
||||
"""
|
||||
import logging
|
||||
|
||||
import numpy as np
|
||||
from sklearn.cluster import DBSCAN
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def cluster_faces(
|
||||
embeddings: np.ndarray,
|
||||
eps: float = 0.35,
|
||||
min_samples: int = 2,
|
||||
) -> np.ndarray:
|
||||
"""Cluster face embeddings using DBSCAN with cosine metric.
|
||||
|
||||
Args:
|
||||
embeddings: (N, D) float32 array of L2-normalized face embeddings.
|
||||
eps: Maximum cosine distance between two samples to be in the
|
||||
same neighborhood. Lower = tighter clusters.
|
||||
min_samples: Minimum cluster size.
|
||||
|
||||
Returns:
|
||||
(N,) int array of cluster labels. -1 = noise / unclustered.
|
||||
"""
|
||||
if len(embeddings) < min_samples:
|
||||
return np.full(len(embeddings), -1, dtype=int)
|
||||
|
||||
db = DBSCAN(eps=eps, min_samples=min_samples, metric="cosine")
|
||||
labels = db.fit_predict(embeddings)
|
||||
|
||||
n_clusters = len(set(labels) - {-1})
|
||||
n_noise = (labels == -1).sum()
|
||||
logger.info(
|
||||
"Face clustering: %d embeddings → %d clusters, %d noise",
|
||||
len(embeddings), n_clusters, n_noise,
|
||||
)
|
||||
return labels
|
||||
@@ -1,138 +0,0 @@
|
||||
"""
|
||||
YOLOv8n object detector using raw ONNX Runtime.
|
||||
|
||||
Expects {models_dir}/detect/yolov8n.onnx, exported from ultralytics
|
||||
via bootstrap_models.py. We do NOT ship ultralytics at runtime to
|
||||
avoid dragging in torch.
|
||||
"""
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
|
||||
from app.config import VisionSettings
|
||||
from app.services.vision.base import ObjectDetector, DetectionBox
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_INPUT_SIZE = 640
|
||||
|
||||
# COCO class names (80 classes)
|
||||
COCO_LABELS = [
|
||||
"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train",
|
||||
"truck", "boat", "traffic light", "fire hydrant", "stop sign",
|
||||
"parking meter", "bench", "bird", "cat", "dog", "horse", "sheep",
|
||||
"cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella",
|
||||
"handbag", "tie", "suitcase", "frisbee", "skis", "snowboard",
|
||||
"sports ball", "kite", "baseball bat", "baseball glove", "skateboard",
|
||||
"surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork",
|
||||
"knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange",
|
||||
"broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair",
|
||||
"couch", "potted plant", "bed", "dining table", "toilet", "tv",
|
||||
"laptop", "mouse", "remote", "keyboard", "cell phone", "microwave",
|
||||
"oven", "toaster", "sink", "refrigerator", "book", "clock", "vase",
|
||||
"scissors", "teddy bear", "hair drier", "toothbrush",
|
||||
]
|
||||
|
||||
|
||||
def _preprocess(image: np.ndarray) -> tuple[np.ndarray, float, float]:
|
||||
"""Letterbox-resize + normalize to NCHW float32. Returns input tensor
|
||||
and scale factors for mapping boxes back to original coords."""
|
||||
from PIL import Image
|
||||
|
||||
img = Image.fromarray(image).convert("RGB")
|
||||
orig_w, orig_h = img.size
|
||||
|
||||
scale = min(_INPUT_SIZE / orig_w, _INPUT_SIZE / orig_h)
|
||||
new_w = int(orig_w * scale)
|
||||
new_h = int(orig_h * scale)
|
||||
img = img.resize((new_w, new_h), Image.BICUBIC)
|
||||
|
||||
# Paste onto gray canvas
|
||||
canvas = np.full((_INPUT_SIZE, _INPUT_SIZE, 3), 114, dtype=np.uint8)
|
||||
pad_x = (_INPUT_SIZE - new_w) // 2
|
||||
pad_y = (_INPUT_SIZE - new_h) // 2
|
||||
canvas[pad_y : pad_y + new_h, pad_x : pad_x + new_w] = np.array(img)
|
||||
|
||||
blob = canvas.astype(np.float32) / 255.0
|
||||
blob = blob.transpose(2, 0, 1)[np.newaxis] # NCHW
|
||||
return blob, scale, pad_x, pad_y
|
||||
|
||||
|
||||
def _postprocess(
|
||||
outputs: np.ndarray,
|
||||
scale: float,
|
||||
pad_x: int,
|
||||
pad_y: int,
|
||||
orig_w: int,
|
||||
orig_h: int,
|
||||
conf_threshold: float,
|
||||
max_detections: int,
|
||||
) -> list[DetectionBox]:
|
||||
"""Parse YOLOv8 output (1, 84, N) → list of DetectionBox."""
|
||||
# outputs shape: (1, 84, N) where 84 = 4 box coords + 80 class scores
|
||||
preds = outputs[0] # (84, N)
|
||||
preds = preds.T # (N, 84)
|
||||
|
||||
boxes_xywh = preds[:, :4]
|
||||
scores = preds[:, 4:]
|
||||
|
||||
class_ids = np.argmax(scores, axis=1)
|
||||
confidences = scores[np.arange(len(scores)), class_ids]
|
||||
|
||||
mask = confidences >= conf_threshold
|
||||
boxes_xywh = boxes_xywh[mask]
|
||||
class_ids = class_ids[mask]
|
||||
confidences = confidences[mask]
|
||||
|
||||
if len(confidences) == 0:
|
||||
return []
|
||||
|
||||
# Sort by confidence, take top N
|
||||
order = np.argsort(-confidences)[:max_detections]
|
||||
boxes_xywh = boxes_xywh[order]
|
||||
class_ids = class_ids[order]
|
||||
confidences = confidences[order]
|
||||
|
||||
results = []
|
||||
for i in range(len(confidences)):
|
||||
cx, cy, w, h = boxes_xywh[i]
|
||||
# Remove letterbox padding and rescale to original image
|
||||
x1 = (cx - w / 2 - pad_x) / scale
|
||||
y1 = (cy - h / 2 - pad_y) / scale
|
||||
x2 = (cx + w / 2 - pad_x) / scale
|
||||
y2 = (cy + h / 2 - pad_y) / scale
|
||||
# Normalize to 0-1
|
||||
bbox = [
|
||||
max(0, x1 / orig_w),
|
||||
max(0, y1 / orig_h),
|
||||
min(1, x2 / orig_w),
|
||||
min(1, y2 / orig_h),
|
||||
]
|
||||
label = COCO_LABELS[class_ids[i]] if class_ids[i] < len(COCO_LABELS) else f"class_{class_ids[i]}"
|
||||
results.append(DetectionBox(label=label, confidence=float(confidences[i]), bbox=bbox))
|
||||
|
||||
return results
|
||||
|
||||
|
||||
class YOLOv8Detector(ObjectDetector):
|
||||
def __init__(self, settings: VisionSettings):
|
||||
model_path = Path(settings.models_dir) / "detect" / "yolov8n.onnx"
|
||||
|
||||
from app.services.vision.providers import create_session
|
||||
|
||||
logger.info("Loading YOLOv8n from %s", model_path)
|
||||
self._session = create_session(str(model_path), configured_providers=settings.execution_providers)
|
||||
self._conf_threshold = settings.detector.min_confidence
|
||||
self._max_detections = settings.detector.max_detections
|
||||
|
||||
def detect(self, image: np.ndarray) -> list[DetectionBox]:
|
||||
orig_h, orig_w = image.shape[:2]
|
||||
blob, scale, pad_x, pad_y = _preprocess(image)
|
||||
input_name = self._session.get_inputs()[0].name
|
||||
outputs = self._session.run(None, {input_name: blob})[0]
|
||||
return _postprocess(
|
||||
outputs, scale, pad_x, pad_y, orig_w, orig_h,
|
||||
self._conf_threshold, self._max_detections,
|
||||
)
|
||||
@@ -1,146 +1,58 @@
|
||||
"""
|
||||
CLIP / SigLIP2 embedder using ONNX Runtime.
|
||||
|
||||
Supports two model families:
|
||||
- OpenCLIP ViT-B/32 (512-d) — legacy, config name "openclip_vitb32"
|
||||
- SigLIP2 ViT-B/16 (768-d) — default, config name "siglip2_vitb16"
|
||||
|
||||
Expects two ONNX files under {models_dir}/embed/:
|
||||
- visual.onnx (image encoder)
|
||||
- textual.onnx (text encoder)
|
||||
|
||||
These are exported from open_clip via export_models.py / bootstrap_models.py.
|
||||
OpenCLIP ViT-B/32 visual encoder (ONNX). Produces 512-d image features
|
||||
consumed by the content classifier. Not exposed as a standalone service;
|
||||
the classifier owns the lifecycle.
|
||||
"""
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
import onnxruntime as ort # noqa: F401 (provider plumbing relies on this)
|
||||
|
||||
from app.config import VisionSettings
|
||||
from app.services.vision.base import Embedder
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ── Model-specific constants ──────────────────────────────────────────
|
||||
|
||||
# OpenCLIP ViT-B/32 (ImageNet norm, 224px)
|
||||
_OPENCLIP_MEAN = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32)
|
||||
_OPENCLIP_STD = np.array([0.26862954, 0.26130258, 0.27577711], dtype=np.float32)
|
||||
_OPENCLIP_SIZE = 224
|
||||
|
||||
# SigLIP2 ViT-B/16 (SigLIP norm, 384px)
|
||||
_SIGLIP2_MEAN = np.array([0.5, 0.5, 0.5], dtype=np.float32)
|
||||
_SIGLIP2_STD = np.array([0.5, 0.5, 0.5], dtype=np.float32)
|
||||
_SIGLIP2_SIZE = 384
|
||||
_MEAN = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32)
|
||||
_STD = np.array([0.26862954, 0.26130258, 0.27577711], dtype=np.float32)
|
||||
_SIZE = 224
|
||||
|
||||
|
||||
def _preprocess_image(
|
||||
image: np.ndarray,
|
||||
input_size: int,
|
||||
mean: np.ndarray,
|
||||
std: np.ndarray,
|
||||
) -> np.ndarray:
|
||||
"""Resize, center-crop, normalize an RGB uint8 image to NCHW float32."""
|
||||
def _preprocess(image: np.ndarray) -> np.ndarray:
|
||||
from PIL import Image
|
||||
|
||||
img = Image.fromarray(image).convert("RGB")
|
||||
w, h = img.size
|
||||
scale = input_size / min(w, h)
|
||||
scale = _SIZE / min(w, h)
|
||||
img = img.resize((int(w * scale), int(h * scale)), Image.BICUBIC)
|
||||
w, h = img.size
|
||||
left = (w - input_size) // 2
|
||||
top = (h - input_size) // 2
|
||||
img = img.crop((left, top, left + input_size, top + input_size))
|
||||
left = (w - _SIZE) // 2
|
||||
top = (h - _SIZE) // 2
|
||||
img = img.crop((left, top, left + _SIZE, top + _SIZE))
|
||||
|
||||
arr = np.array(img, dtype=np.float32) / 255.0
|
||||
arr = (arr - mean) / std
|
||||
arr = arr.transpose(2, 0, 1) # HWC → CHW
|
||||
return arr[np.newaxis] # NCHW
|
||||
arr = (arr - _MEAN) / _STD
|
||||
arr = arr.transpose(2, 0, 1)
|
||||
return arr[np.newaxis]
|
||||
|
||||
|
||||
class OpenCLIPEmbedder(Embedder):
|
||||
"""Legacy OpenCLIP ViT-B/32 embedder (512-d)."""
|
||||
class CLIPVisualEncoder:
|
||||
"""OpenCLIP ViT-B/32 image encoder, 512-d normalized output."""
|
||||
|
||||
def __init__(self, settings: VisionSettings):
|
||||
model_dir = Path(settings.models_dir) / "embed"
|
||||
visual_path = model_dir / "visual.onnx"
|
||||
textual_path = model_dir / "textual.onnx"
|
||||
|
||||
model_path = Path(settings.models_dir) / "embed" / "visual.onnx"
|
||||
from app.services.vision.providers import create_session
|
||||
from app.config import settings as app_settings
|
||||
providers = app_settings.vision.execution_providers
|
||||
|
||||
logger.info("Loading OpenCLIP visual encoder from %s", visual_path)
|
||||
self._visual = create_session(str(visual_path), configured_providers=providers)
|
||||
logger.info("Loading CLIP visual encoder from %s", model_path)
|
||||
self._session = create_session(
|
||||
str(model_path),
|
||||
configured_providers=app_settings.vision.execution_providers,
|
||||
)
|
||||
|
||||
logger.info("Loading OpenCLIP textual encoder from %s", textual_path)
|
||||
self._textual = create_session(str(textual_path), configured_providers=providers)
|
||||
|
||||
def embed_image(self, image: np.ndarray) -> np.ndarray:
|
||||
inp = _preprocess_image(image, _OPENCLIP_SIZE, _OPENCLIP_MEAN, _OPENCLIP_STD)
|
||||
input_name = self._visual.get_inputs()[0].name
|
||||
out = self._visual.run(None, {input_name: inp})[0][0]
|
||||
def encode(self, image: np.ndarray) -> np.ndarray:
|
||||
inp = _preprocess(image)
|
||||
name = self._session.get_inputs()[0].name
|
||||
out = self._session.run(None, {name: inp})[0][0]
|
||||
out = out / np.linalg.norm(out)
|
||||
return out.astype(np.float32)
|
||||
|
||||
def embed_text(self, text: str) -> np.ndarray:
|
||||
import open_clip
|
||||
tokenizer = open_clip.get_tokenizer("ViT-B-32")
|
||||
tokens = tokenizer([text]).numpy().astype(np.int64)
|
||||
eot_indices = tokens.argmax(axis=-1).astype(np.int64)
|
||||
inputs = self._textual.get_inputs()
|
||||
out = self._textual.run(None, {
|
||||
inputs[0].name: tokens,
|
||||
inputs[1].name: eot_indices,
|
||||
})[0][0]
|
||||
out = out / np.linalg.norm(out)
|
||||
return out.astype(np.float32)
|
||||
|
||||
@property
|
||||
def dim(self) -> int:
|
||||
return 512
|
||||
|
||||
|
||||
class SigLIP2Embedder(Embedder):
|
||||
"""SigLIP2 ViT-B/16 embedder (768-d) — higher recall than OpenCLIP."""
|
||||
|
||||
def __init__(self, settings: VisionSettings):
|
||||
model_dir = Path(settings.models_dir) / "embed_siglip2"
|
||||
visual_path = model_dir / "visual.onnx"
|
||||
textual_path = model_dir / "textual.onnx"
|
||||
|
||||
from app.services.vision.providers import create_session
|
||||
from app.config import settings as app_settings
|
||||
providers = app_settings.vision.execution_providers
|
||||
|
||||
logger.info("Loading SigLIP2 visual encoder from %s", visual_path)
|
||||
self._visual = create_session(str(visual_path), configured_providers=providers)
|
||||
|
||||
logger.info("Loading SigLIP2 textual encoder from %s", textual_path)
|
||||
self._textual = create_session(str(textual_path), configured_providers=providers)
|
||||
|
||||
def embed_image(self, image: np.ndarray) -> np.ndarray:
|
||||
inp = _preprocess_image(image, _SIGLIP2_SIZE, _SIGLIP2_MEAN, _SIGLIP2_STD)
|
||||
input_name = self._visual.get_inputs()[0].name
|
||||
out = self._visual.run(None, {input_name: inp})[0][0]
|
||||
out = out / np.linalg.norm(out)
|
||||
return out.astype(np.float32)
|
||||
|
||||
def embed_text(self, text: str) -> np.ndarray:
|
||||
import open_clip
|
||||
tokenizer = open_clip.get_tokenizer("ViT-B-16-SigLIP-384")
|
||||
tokens = tokenizer([text]).numpy().astype(np.int64)
|
||||
inputs = self._textual.get_inputs()
|
||||
feed = {inputs[0].name: tokens}
|
||||
# SigLIP2 text encoder may need attention mask
|
||||
if len(inputs) > 1:
|
||||
attention_mask = (tokens != 0).astype(np.int64)
|
||||
feed[inputs[1].name] = attention_mask
|
||||
out = self._textual.run(None, feed)[0][0]
|
||||
out = out / np.linalg.norm(out)
|
||||
return out.astype(np.float32)
|
||||
|
||||
@property
|
||||
def dim(self) -> int:
|
||||
return 768
|
||||
|
||||
@@ -1,253 +1,61 @@
|
||||
"""
|
||||
Export / download all vision model weights to ONNX format.
|
||||
Export the OpenCLIP ViT-B/32 visual encoder to ONNX.
|
||||
|
||||
Run ONCE on any machine with Python + pip (doesn't need GPU):
|
||||
Run once on any machine with Python + pip (no GPU needed):
|
||||
|
||||
pip install open-clip-torch ultralytics onnx
|
||||
pip install open-clip-torch onnx
|
||||
python -m app.services.vision.export_models [--models-dir /data/models]
|
||||
|
||||
This produces:
|
||||
embed/visual.onnx (~350 MB)
|
||||
embed/textual.onnx (~250 MB)
|
||||
detect/yolov8n.onnx (~12 MB)
|
||||
|
||||
YuNet and SFace are downloaded by bootstrap_models.py at worker boot
|
||||
(Apache 2.0, lightweight, no export step needed).
|
||||
|
||||
After export, copy the /data/models directory into your Docker volume:
|
||||
docker cp /data/models mulita-worker:/data/models
|
||||
Or mount a host path in docker-compose.yml.
|
||||
Produces:
|
||||
embed/visual.onnx (~350 MB)
|
||||
"""
|
||||
import argparse
|
||||
import logging
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def export_openclip(models_dir: Path):
|
||||
"""Export OpenCLIP ViT-B/32 to two ONNX files (visual + textual)."""
|
||||
def export_openclip_visual(models_dir: Path):
|
||||
import torch
|
||||
import open_clip
|
||||
|
||||
out_dir = models_dir / "embed"
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
visual_path = out_dir / "visual.onnx"
|
||||
textual_path = out_dir / "textual.onnx"
|
||||
|
||||
if visual_path.exists() and textual_path.exists():
|
||||
logger.info("OpenCLIP ONNX files already exist, skipping export")
|
||||
if visual_path.exists():
|
||||
logger.info("OpenCLIP visual.onnx already exists, skipping export")
|
||||
return
|
||||
|
||||
logger.info("Loading OpenCLIP ViT-B-32 laion2b_s34b_b79k...")
|
||||
model, _, preprocess = open_clip.create_model_and_transforms(
|
||||
model, _, _ = open_clip.create_model_and_transforms(
|
||||
"ViT-B-32", pretrained="laion2b_s34b_b79k"
|
||||
)
|
||||
model.eval()
|
||||
|
||||
# Use dynamo=False to get the legacy TorchScript exporter which
|
||||
# produces IR version 9 (compatible with onnxruntime 1.17.x).
|
||||
# The new torch.onnx.export default (dynamo=True) emits IR 10.
|
||||
export_kwargs = dict(opset_version=14, dynamo=False)
|
||||
|
||||
# ── Visual encoder ────────────────────────────────────────────────
|
||||
if not visual_path.exists():
|
||||
logger.info("Exporting visual encoder → %s", visual_path)
|
||||
dummy_image = torch.randn(1, 3, 224, 224)
|
||||
torch.onnx.export(
|
||||
model.visual,
|
||||
dummy_image,
|
||||
str(visual_path),
|
||||
input_names=["image"],
|
||||
output_names=["embedding"],
|
||||
dynamic_axes={"image": {0: "batch"}},
|
||||
**export_kwargs,
|
||||
)
|
||||
size_mb = visual_path.stat().st_size / 1e6
|
||||
logger.info("Visual encoder exported (%.1f MB)", size_mb)
|
||||
|
||||
# ── Textual encoder ───────────────────────────────────────────────
|
||||
if not textual_path.exists():
|
||||
logger.info("Exporting textual encoder → %s", textual_path)
|
||||
tokenizer = open_clip.get_tokenizer("ViT-B-32")
|
||||
dummy_text = tokenizer(["a photo"]).to(torch.int64)
|
||||
|
||||
class TextEncoder(torch.nn.Module):
|
||||
"""Wrap the CLIP text encoder to avoid argmax in the ONNX graph.
|
||||
OpenCLIP uses argmax to find the EOT token position, but ORT
|
||||
ARM64 doesn't support ArgMax(13). We pre-compute the EOT index
|
||||
from the token sequence and pass it directly."""
|
||||
def __init__(self, clip_model):
|
||||
super().__init__()
|
||||
self.transformer = clip_model.transformer
|
||||
self.token_embedding = clip_model.token_embedding
|
||||
self.positional_embedding = clip_model.positional_embedding
|
||||
self.ln_final = clip_model.ln_final
|
||||
self.text_projection = clip_model.text_projection
|
||||
|
||||
def forward(self, text, eot_indices):
|
||||
x = self.token_embedding(text)
|
||||
x = x + self.positional_embedding
|
||||
x = x.permute(1, 0, 2) # NLD -> LND
|
||||
x = self.transformer(x)
|
||||
x = x.permute(1, 0, 2) # LND -> NLD
|
||||
x = self.ln_final(x)
|
||||
# Take the feature at the EOT token. The EOT index is
|
||||
# passed in as a separate input (computed outside ONNX)
|
||||
# to avoid ArgMax(13) which ORT ARM64 doesn't support.
|
||||
x = x[torch.arange(x.shape[0]), eot_indices]
|
||||
x = x @ self.text_projection
|
||||
return x
|
||||
|
||||
text_enc = TextEncoder(model)
|
||||
text_enc.eval()
|
||||
|
||||
# Compute EOT indices from dummy tokens (argmax of token ids)
|
||||
dummy_eot = dummy_text.argmax(dim=-1)
|
||||
|
||||
torch.onnx.export(
|
||||
text_enc,
|
||||
(dummy_text, dummy_eot),
|
||||
str(textual_path),
|
||||
input_names=["text", "eot_indices"],
|
||||
output_names=["embedding"],
|
||||
dynamic_axes={"text": {0: "batch"}, "eot_indices": {0: "batch"}},
|
||||
**export_kwargs,
|
||||
)
|
||||
size_mb = textual_path.stat().st_size / 1e6
|
||||
logger.info("Textual encoder exported (%.1f MB)", size_mb)
|
||||
|
||||
|
||||
def export_siglip2(models_dir: Path):
|
||||
"""Export SigLIP2 ViT-B/16 to two ONNX files (visual + textual)."""
|
||||
import torch
|
||||
import open_clip
|
||||
|
||||
out_dir = models_dir / "embed_siglip2"
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
visual_path = out_dir / "visual.onnx"
|
||||
textual_path = out_dir / "textual.onnx"
|
||||
|
||||
if visual_path.exists() and textual_path.exists():
|
||||
logger.info("SigLIP2 ONNX files already exist, skipping export")
|
||||
return
|
||||
|
||||
logger.info("Loading SigLIP2 ViT-B-16-SigLIP-384 webli...")
|
||||
model, _, preprocess = open_clip.create_model_and_transforms(
|
||||
"ViT-B-16-SigLIP-384", pretrained="webli"
|
||||
logger.info("Exporting visual encoder → %s", visual_path)
|
||||
dummy = torch.randn(1, 3, 224, 224)
|
||||
torch.onnx.export(
|
||||
model.visual,
|
||||
dummy,
|
||||
str(visual_path),
|
||||
input_names=["image"],
|
||||
output_names=["embedding"],
|
||||
dynamic_axes={"image": {0: "batch"}},
|
||||
opset_version=14,
|
||||
dynamo=False,
|
||||
)
|
||||
model.eval()
|
||||
|
||||
export_kwargs = dict(opset_version=14, dynamo=False)
|
||||
|
||||
# ── Visual encoder ────────────────────────────────────────────────
|
||||
if not visual_path.exists():
|
||||
logger.info("Exporting SigLIP2 visual encoder → %s", visual_path)
|
||||
dummy_image = torch.randn(1, 3, 384, 384)
|
||||
torch.onnx.export(
|
||||
model.visual,
|
||||
dummy_image,
|
||||
str(visual_path),
|
||||
input_names=["image"],
|
||||
output_names=["embedding"],
|
||||
dynamic_axes={"image": {0: "batch"}},
|
||||
**export_kwargs,
|
||||
)
|
||||
size_mb = visual_path.stat().st_size / 1e6
|
||||
logger.info("SigLIP2 visual encoder exported (%.1f MB)", size_mb)
|
||||
|
||||
# ── Textual encoder ───────────────────────────────────────────────
|
||||
if not textual_path.exists():
|
||||
logger.info("Exporting SigLIP2 textual encoder → %s", textual_path)
|
||||
tokenizer = open_clip.get_tokenizer("ViT-B-16-SigLIP-384")
|
||||
dummy_text = tokenizer(["a photo"]).to(torch.int64)
|
||||
|
||||
class SigLIP2TextEncoder(torch.nn.Module):
|
||||
"""Wrap the SigLIP2 text transformer for ONNX export."""
|
||||
def __init__(self, clip_model):
|
||||
super().__init__()
|
||||
self.text = clip_model.text
|
||||
|
||||
def forward(self, text):
|
||||
return self.text(text)
|
||||
|
||||
text_enc = SigLIP2TextEncoder(model)
|
||||
text_enc.eval()
|
||||
|
||||
torch.onnx.export(
|
||||
text_enc,
|
||||
dummy_text,
|
||||
str(textual_path),
|
||||
input_names=["text"],
|
||||
output_names=["embedding"],
|
||||
dynamic_axes={"text": {0: "batch"}},
|
||||
**export_kwargs,
|
||||
)
|
||||
size_mb = textual_path.stat().st_size / 1e6
|
||||
logger.info("SigLIP2 textual encoder exported (%.1f MB)", size_mb)
|
||||
|
||||
|
||||
def export_yolov8n(models_dir: Path):
|
||||
"""Export YOLOv8n to ONNX."""
|
||||
out_dir = models_dir / "detect"
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
onnx_path = out_dir / "yolov8n.onnx"
|
||||
|
||||
if onnx_path.exists():
|
||||
logger.info("YOLOv8n ONNX already exists, skipping export")
|
||||
return
|
||||
|
||||
logger.info("Exporting YOLOv8n → %s", onnx_path)
|
||||
|
||||
from ultralytics import YOLO
|
||||
|
||||
model = YOLO("yolov8n.pt")
|
||||
model.export(format="onnx", imgsz=640, simplify=True)
|
||||
|
||||
# ultralytics exports to cwd as yolov8n.onnx — move to target. Use
|
||||
# shutil.move rather than Path.rename so it works across filesystems
|
||||
# (the cwd is typically /app inside the container, while the target
|
||||
# /data/models is a separately-mounted volume — Path.rename raises
|
||||
# "Invalid cross-device link" in that case).
|
||||
import shutil
|
||||
|
||||
exported = Path("yolov8n.onnx")
|
||||
if exported.exists():
|
||||
shutil.move(str(exported), str(onnx_path))
|
||||
|
||||
size_mb = onnx_path.stat().st_size / 1e6
|
||||
logger.info("YOLOv8n exported (%.1f MB)", size_mb)
|
||||
logger.info("Visual encoder exported (%.1f MB)", visual_path.stat().st_size / 1e6)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Export vision model weights to ONNX")
|
||||
parser.add_argument(
|
||||
"--models-dir",
|
||||
type=Path,
|
||||
default=Path("/data/models"),
|
||||
help="Directory to write model files (default: /data/models)",
|
||||
)
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--models-dir", type=Path, default=Path("/data/models"))
|
||||
args = parser.parse_args()
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s %(levelname)s %(message)s",
|
||||
)
|
||||
|
||||
models_dir = args.models_dir
|
||||
models_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
logger.info("Exporting models to %s", models_dir)
|
||||
|
||||
export_openclip(models_dir)
|
||||
export_siglip2(models_dir)
|
||||
export_yolov8n(models_dir)
|
||||
|
||||
logger.info("Done. Run bootstrap_models.py next to download YuNet + SFace.")
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
args.models_dir.mkdir(parents=True, exist_ok=True)
|
||||
export_openclip_visual(args.models_dir)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -1,146 +0,0 @@
|
||||
"""
|
||||
Face detection (YuNet) + recognition (SFace) using OpenCV DNN.
|
||||
|
||||
YuNet is loaded via cv2.FaceDetectorYN which handles the multi-scale
|
||||
anchor decoding and NMS internally. SFace recognition uses raw ONNX
|
||||
Runtime for the 128-d embedding.
|
||||
|
||||
Both models are from opencv_zoo (Apache 2.0 license).
|
||||
Expects {models_dir}/face/:
|
||||
- yunet.onnx (~233 KB)
|
||||
- sface.onnx (~37 MB, 128-d embeddings)
|
||||
"""
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import cv2
|
||||
import onnxruntime as ort
|
||||
|
||||
from app.config import VisionSettings
|
||||
from app.services.vision.base import FaceProcessor, FaceDetection
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _align_face(image: np.ndarray, landmarks: np.ndarray) -> np.ndarray:
|
||||
"""Align and crop a 112x112 face patch using 5-point landmarks."""
|
||||
left_eye = landmarks[0]
|
||||
right_eye = landmarks[1]
|
||||
|
||||
dx = right_eye[0] - left_eye[0]
|
||||
dy = right_eye[1] - left_eye[1]
|
||||
angle = np.degrees(np.arctan2(dy, dx))
|
||||
eye_center = ((left_eye[0] + right_eye[0]) / 2, (left_eye[1] + right_eye[1]) / 2)
|
||||
eye_dist = np.sqrt(dx * dx + dy * dy)
|
||||
|
||||
M = cv2.getRotationMatrix2D(eye_center, angle, 1.0)
|
||||
rotated = cv2.warpAffine(image, M, (image.shape[1], image.shape[0]))
|
||||
|
||||
# Crop around face center
|
||||
scale = 64.0 / max(eye_dist, 1e-6)
|
||||
cx, cy = eye_center
|
||||
half = 56.0 / scale
|
||||
x1 = max(0, int(cx - half))
|
||||
y1 = max(0, int(cy - half * 0.8))
|
||||
x2 = min(rotated.shape[1], int(cx + half))
|
||||
y2 = min(rotated.shape[0], int(cy + half * 1.2))
|
||||
crop = rotated[y1:y2, x1:x2]
|
||||
|
||||
if crop.size == 0:
|
||||
return np.zeros((112, 112, 3), dtype=np.float32)
|
||||
|
||||
return cv2.resize(crop, (112, 112)).astype(np.float32)
|
||||
|
||||
|
||||
class YuNetSFaceProcessor(FaceProcessor):
|
||||
def __init__(self, settings: VisionSettings):
|
||||
face_dir = Path(settings.models_dir) / "face"
|
||||
yunet_path = str(face_dir / "yunet.onnx")
|
||||
sface_path = str(face_dir / "sface.onnx")
|
||||
|
||||
# YuNet via OpenCV's FaceDetectorYN — handles anchor decoding + NMS
|
||||
self._detector = cv2.FaceDetectorYN.create(
|
||||
yunet_path,
|
||||
"",
|
||||
(640, 640),
|
||||
settings.faces.recognition_threshold,
|
||||
0.3, # NMS threshold
|
||||
5000, # top_k
|
||||
)
|
||||
logger.info("YuNet face detector loaded via OpenCV")
|
||||
|
||||
# SFace via ONNX Runtime
|
||||
from app.services.vision.providers import create_session
|
||||
ort.set_default_logger_severity(3)
|
||||
self._recognizer = create_session(sface_path, configured_providers=settings.execution_providers)
|
||||
logger.info("SFace recognizer loaded via ONNX Runtime")
|
||||
|
||||
self._min_face_size = settings.faces.min_face_size
|
||||
|
||||
def process(self, image: np.ndarray) -> list[FaceDetection]:
|
||||
orig_h, orig_w = image.shape[:2]
|
||||
|
||||
# Convert RGB → BGR for OpenCV
|
||||
bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
|
||||
|
||||
# Set input size to actual image dimensions
|
||||
self._detector.setInputSize((orig_w, orig_h))
|
||||
|
||||
# Detect faces
|
||||
_, faces_raw = self._detector.detect(bgr)
|
||||
|
||||
if faces_raw is None or len(faces_raw) == 0:
|
||||
return []
|
||||
|
||||
results = []
|
||||
for face in faces_raw:
|
||||
# face: [x, y, w, h, right_eye_x, right_eye_y, left_eye_x, left_eye_y,
|
||||
# nose_x, nose_y, right_mouth_x, right_mouth_y, left_mouth_x, left_mouth_y, score]
|
||||
x, y, w, h = int(face[0]), int(face[1]), int(face[2]), int(face[3])
|
||||
score = float(face[14])
|
||||
|
||||
# Filter small faces
|
||||
face_size = max(w, h)
|
||||
if face_size < self._min_face_size:
|
||||
continue
|
||||
|
||||
# Normalized bbox
|
||||
bbox = [
|
||||
max(0, x / orig_w),
|
||||
max(0, y / orig_h),
|
||||
min(1, (x + w) / orig_w),
|
||||
min(1, (y + h) / orig_h),
|
||||
]
|
||||
|
||||
# Extract 5-point landmarks for alignment
|
||||
landmarks = np.array([
|
||||
[face[4], face[5]], # right eye
|
||||
[face[6], face[7]], # left eye
|
||||
[face[8], face[9]], # nose
|
||||
[face[10], face[11]], # right mouth
|
||||
[face[12], face[13]], # left mouth
|
||||
], dtype=np.float32)
|
||||
|
||||
# Align face for recognition
|
||||
face_crop = _align_face(image, landmarks)
|
||||
|
||||
# SFace expects (1, 3, 112, 112) float32, BGR
|
||||
face_bgr = cv2.cvtColor(face_crop.astype(np.uint8), cv2.COLOR_RGB2BGR)
|
||||
face_blob = (face_bgr.astype(np.float32) / 255.0).transpose(2, 0, 1)[np.newaxis]
|
||||
|
||||
rec_input = self._recognizer.get_inputs()[0].name
|
||||
embedding = self._recognizer.run(None, {rec_input: face_blob})[0][0]
|
||||
embedding = embedding / np.linalg.norm(embedding)
|
||||
|
||||
results.append(FaceDetection(
|
||||
bbox=bbox,
|
||||
embedding=embedding.astype(np.float32),
|
||||
quality=score,
|
||||
))
|
||||
|
||||
return results
|
||||
|
||||
@property
|
||||
def embedding_dim(self) -> int:
|
||||
return 128
|
||||
@@ -1,73 +0,0 @@
|
||||
"""
|
||||
Face detection + recognition using InsightFace (RetinaFace + ArcFace).
|
||||
|
||||
Uses the buffalo_l model pack which auto-downloads on first use (~300MB).
|
||||
Produces 512-d ArcFace embeddings. Non-commercial research license —
|
||||
fine for homelab self-hosting.
|
||||
"""
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
from app.config import VisionSettings
|
||||
from app.services.vision.base import FaceProcessor, FaceDetection
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class InsightFaceProcessor(FaceProcessor):
|
||||
def __init__(self, settings: VisionSettings):
|
||||
from insightface.app import FaceAnalysis
|
||||
|
||||
model_root = str(Path(settings.models_dir) / "face" / "insightface")
|
||||
logger.info("Loading InsightFace buffalo_l from %s", model_root)
|
||||
|
||||
from app.services.vision.providers import get_providers
|
||||
providers = get_providers(settings.execution_providers)
|
||||
|
||||
self._app = FaceAnalysis(
|
||||
name="buffalo_l",
|
||||
root=model_root,
|
||||
providers=providers,
|
||||
)
|
||||
self._app.prepare(ctx_id=-1, det_size=(640, 640))
|
||||
self._min_det_score = settings.faces.recognition_threshold
|
||||
|
||||
def process(self, image: np.ndarray) -> list[FaceDetection]:
|
||||
orig_h, orig_w = image.shape[:2]
|
||||
|
||||
# InsightFace expects BGR
|
||||
bgr = image[:, :, ::-1].copy()
|
||||
|
||||
faces = self._app.get(bgr)
|
||||
|
||||
if not faces:
|
||||
return []
|
||||
|
||||
results = []
|
||||
for face in faces:
|
||||
if face.det_score < self._min_det_score:
|
||||
continue
|
||||
|
||||
# face.bbox is [x1, y1, x2, y2] in pixel coords
|
||||
x1, y1, x2, y2 = face.bbox
|
||||
bbox = [
|
||||
max(0, float(x1) / orig_w),
|
||||
max(0, float(y1) / orig_h),
|
||||
min(1, float(x2) / orig_w),
|
||||
min(1, float(y2) / orig_h),
|
||||
]
|
||||
|
||||
embedding = face.normed_embedding # already L2-normalized, 512-d
|
||||
results.append(FaceDetection(
|
||||
bbox=bbox,
|
||||
embedding=embedding.astype(np.float32),
|
||||
quality=float(face.det_score),
|
||||
))
|
||||
|
||||
return results
|
||||
|
||||
@property
|
||||
def embedding_dim(self) -> int:
|
||||
return 512
|
||||
@@ -1,45 +0,0 @@
|
||||
"""
|
||||
OCR engine using rapidocr-onnxruntime (PP-OCRv4 weights).
|
||||
|
||||
No PaddlePaddle dependency — pure ONNX Runtime. Language packs are
|
||||
downloaded automatically by rapidocr on first use.
|
||||
"""
|
||||
import logging
|
||||
|
||||
import numpy as np
|
||||
|
||||
from app.config import VisionSettings
|
||||
from app.services.vision.base import OCREngine, OCRResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RapidOCREngine(OCREngine):
|
||||
def __init__(self, settings: VisionSettings):
|
||||
from rapidocr_onnxruntime import RapidOCR
|
||||
|
||||
self._min_confidence = settings.ocr.min_confidence
|
||||
self._engine = RapidOCR()
|
||||
logger.info("RapidOCR engine initialized")
|
||||
|
||||
def run(self, image: np.ndarray) -> list[OCRResult]:
|
||||
result, _ = self._engine(image)
|
||||
if not result:
|
||||
return []
|
||||
|
||||
out = []
|
||||
for box, text, score in result:
|
||||
if score < self._min_confidence:
|
||||
continue
|
||||
# box is [[x1,y1],[x2,y2],[x3,y3],[x4,y4]] — take bounding rect
|
||||
xs = [p[0] for p in box]
|
||||
ys = [p[1] for p in box]
|
||||
h, w = image.shape[:2]
|
||||
bbox = [
|
||||
min(xs) / w,
|
||||
min(ys) / h,
|
||||
max(xs) / w,
|
||||
max(ys) / h,
|
||||
]
|
||||
out.append(OCRResult(text=text, confidence=float(score), bbox=bbox))
|
||||
return out
|
||||
@@ -1,46 +0,0 @@
|
||||
"""
|
||||
ONNX Runtime backend — default CPU inference for all vision models.
|
||||
|
||||
Each create_* method returns a concrete implementation of the
|
||||
corresponding ABC from base.py. Models are loaded from ONNX files
|
||||
under settings.vision.models_dir, downloaded on first boot by
|
||||
bootstrap_models.py.
|
||||
"""
|
||||
import logging
|
||||
|
||||
from app.config import VisionSettings
|
||||
from app.services.vision.base import Embedder, OCREngine, ObjectDetector, FaceProcessor, ContentClassifier
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ONNXBackend:
|
||||
"""Factory for ONNX Runtime-based vision model instances."""
|
||||
|
||||
def __init__(self, vision_settings: VisionSettings):
|
||||
self._settings = vision_settings
|
||||
|
||||
def create_embedder(self) -> Embedder:
|
||||
model_name = self._settings.embedder.name
|
||||
if model_name.startswith("siglip2"):
|
||||
from app.services.vision.embed import SigLIP2Embedder
|
||||
return SigLIP2Embedder(self._settings)
|
||||
else:
|
||||
from app.services.vision.embed import OpenCLIPEmbedder
|
||||
return OpenCLIPEmbedder(self._settings)
|
||||
|
||||
def create_ocr(self) -> OCREngine:
|
||||
from app.services.vision.ocr import RapidOCREngine
|
||||
return RapidOCREngine(self._settings)
|
||||
|
||||
def create_detector(self) -> ObjectDetector:
|
||||
from app.services.vision.detect import YOLOv8Detector
|
||||
return YOLOv8Detector(self._settings)
|
||||
|
||||
def create_face_processor(self) -> FaceProcessor:
|
||||
from app.services.vision.insightface_processor import InsightFaceProcessor
|
||||
return InsightFaceProcessor(self._settings)
|
||||
|
||||
def create_classifier(self) -> ContentClassifier:
|
||||
from app.services.vision.classify import CLIPContentClassifier
|
||||
return CLIPContentClassifier(self._settings)
|
||||
@@ -1,84 +1,29 @@
|
||||
"""
|
||||
ModelRegistry — singleton that lazy-loads vision models per worker process.
|
||||
|
||||
Usage from Celery tasks:
|
||||
|
||||
from app.services.vision.registry import registry
|
||||
embedder = registry.get_embedder()
|
||||
vec = embedder.embed_image(img)
|
||||
|
||||
Models are created on first access and cached for the worker's lifetime.
|
||||
The registry reads settings.vision to decide which backend to use and
|
||||
where model weights live.
|
||||
ModelRegistry — lazy-loads the single content classifier per worker.
|
||||
"""
|
||||
import logging
|
||||
from functools import lru_cache
|
||||
|
||||
from app.config import settings
|
||||
from app.services.vision.base import Embedder, OCREngine, ObjectDetector, FaceProcessor, ContentClassifier
|
||||
from app.services.vision.base import ContentClassifier
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ModelRegistry:
|
||||
"""Central access point for all vision models."""
|
||||
|
||||
def __init__(self):
|
||||
self._vision = settings.vision
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_embedder(self) -> Embedder:
|
||||
logger.info("Loading embedder: %s (backend=%s)", self._vision.embedder.name, self._vision.backend)
|
||||
return self._load_backend().create_embedder()
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_ocr(self) -> OCREngine:
|
||||
logger.info("Loading OCR engine (backend=%s)", self._vision.backend)
|
||||
return self._load_backend().create_ocr()
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_detector(self) -> ObjectDetector:
|
||||
logger.info("Loading object detector (backend=%s)", self._vision.backend)
|
||||
return self._load_backend().create_detector()
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_face_processor(self) -> FaceProcessor:
|
||||
logger.info("Loading face processor (backend=%s)", self._vision.backend)
|
||||
return self._load_backend().create_face_processor()
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_classifier(self) -> ContentClassifier:
|
||||
logger.info("Loading content classifier (backend=%s)", self._vision.backend)
|
||||
return self._load_backend().create_classifier()
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _load_backend(self):
|
||||
"""Import and instantiate the configured backend."""
|
||||
backend_name = self._vision.backend
|
||||
if backend_name == "onnx":
|
||||
from app.services.vision.onnx_backend import ONNXBackend
|
||||
return ONNXBackend(self._vision)
|
||||
elif backend_name == "rocm":
|
||||
from app.services.vision.rocm_backend import ROCmBackend
|
||||
return ROCmBackend(self._vision)
|
||||
else:
|
||||
raise ValueError(f"Unknown vision backend: {backend_name}")
|
||||
from app.services.vision.classify import CLIPContentClassifier
|
||||
return CLIPContentClassifier(self._vision)
|
||||
|
||||
def warmup(self):
|
||||
"""Pre-load all enabled models. Called from Celery worker_process_init
|
||||
on the vision queue to avoid cold-start latency on the first task."""
|
||||
logger.info("Warming up vision models...")
|
||||
self.get_embedder()
|
||||
if self._vision.ocr.enabled:
|
||||
self.get_ocr()
|
||||
if self._vision.detector.enabled:
|
||||
self.get_detector()
|
||||
if self._vision.faces.enabled:
|
||||
self.get_face_processor()
|
||||
if self._vision.classifier.enabled:
|
||||
self.get_classifier()
|
||||
logger.info("Vision model warmup complete")
|
||||
logger.info("Warming up vision classifier...")
|
||||
self.get_classifier()
|
||||
logger.info("Vision warmup complete")
|
||||
|
||||
|
||||
# Module-level singleton. Import this from tasks.
|
||||
registry = ModelRegistry()
|
||||
|
||||
@@ -1,25 +0,0 @@
|
||||
"""
|
||||
ROCm backend — GPU-accelerated inference for Radeon 760M-class hardware.
|
||||
|
||||
Stub: raises NotImplementedError on all factory methods. To enable,
|
||||
set `vision.backend: rocm` in mulita.yml once ROCm support is implemented.
|
||||
"""
|
||||
from app.config import VisionSettings
|
||||
from app.services.vision.base import Embedder, OCREngine, ObjectDetector, FaceProcessor
|
||||
|
||||
|
||||
class ROCmBackend:
|
||||
def __init__(self, vision_settings: VisionSettings):
|
||||
self._settings = vision_settings
|
||||
|
||||
def create_embedder(self) -> Embedder:
|
||||
raise NotImplementedError("ROCm backend not yet implemented — use 'onnx'")
|
||||
|
||||
def create_ocr(self) -> OCREngine:
|
||||
raise NotImplementedError("ROCm backend not yet implemented — use 'onnx'")
|
||||
|
||||
def create_detector(self) -> ObjectDetector:
|
||||
raise NotImplementedError("ROCm backend not yet implemented — use 'onnx'")
|
||||
|
||||
def create_face_processor(self) -> FaceProcessor:
|
||||
raise NotImplementedError("ROCm backend not yet implemented — use 'onnx'")
|
||||
Reference in New Issue
Block a user