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:
@@ -1,33 +1,21 @@
|
||||
"""
|
||||
Celery tasks for the vision pipeline — embedding, OCR, object detection,
|
||||
face recognition.
|
||||
Celery tasks for the vision pipeline.
|
||||
|
||||
All tasks run on the dedicated `vision` queue with limited concurrency
|
||||
(memory-bound CPU inference). They read thumbnails generated by
|
||||
generate_thumbnails, so they MUST run after thumbs complete.
|
||||
|
||||
DB access uses sync psycopg2 sessions (not asyncpg) because Celery
|
||||
forks workers and asyncpg connections can't be shared across forks.
|
||||
A single binary classifier decides whether a photo is 'photography' or
|
||||
'other'. Photos classified as 'other' get needs_review=true so the user
|
||||
can triage screenshots / documents / memes in the UI.
|
||||
"""
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
from celery import shared_task
|
||||
from sqlalchemy import create_engine, text as sa_text, select, delete
|
||||
from sqlalchemy import create_engine, text as sa_text, select, delete, update
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
from PIL import Image
|
||||
|
||||
from app.models.embeddings import Embedding
|
||||
from app.config import settings
|
||||
from app.services.feature_flags import (
|
||||
is_enabled,
|
||||
FLAG_VISION_ENABLED,
|
||||
FLAG_OCR_ENABLED,
|
||||
FLAG_DETECTOR_ENABLED,
|
||||
FLAG_FACES_ENABLED,
|
||||
FLAG_CLASSIFIER_ENABLED,
|
||||
)
|
||||
from app.services.feature_flags import is_enabled, FLAG_VISION_ENABLED
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -35,7 +23,6 @@ VISION_READY_KEY = "mulita:vision:ready"
|
||||
|
||||
|
||||
def _vision_worker_ready() -> bool:
|
||||
"""Check whether the vision worker has finished model bootstrap."""
|
||||
try:
|
||||
import redis as _redis
|
||||
return bool(_redis.from_url(settings.redis_url).exists(VISION_READY_KEY))
|
||||
@@ -47,7 +34,6 @@ _sync_engine = None
|
||||
|
||||
|
||||
def _get_sync_engine():
|
||||
"""Return a module-level singleton engine (one per worker process)."""
|
||||
global _sync_engine
|
||||
if _sync_engine is None:
|
||||
sync_url = settings.database_url.replace("+asyncpg", "+psycopg2").replace("+aiosqlite", "")
|
||||
@@ -56,21 +42,13 @@ def _get_sync_engine():
|
||||
|
||||
|
||||
def _get_sync_session() -> Session:
|
||||
"""Create a sync DB session backed by the shared engine."""
|
||||
return sessionmaker(bind=_get_sync_engine())()
|
||||
|
||||
|
||||
def _load_thumb(photo_id: str, size: str = "medium") -> np.ndarray | None:
|
||||
"""Load a thumbnail as an RGB numpy array.
|
||||
|
||||
Thumbnails may live at ``/data/thumbs/{photo_id}/`` (legacy) or
|
||||
``/data/thumbs/{user_id}/{photo_id}/`` (multi-user). Try both.
|
||||
"""
|
||||
thumb_base = Path("/data/thumbs")
|
||||
# Try legacy flat path first.
|
||||
thumb_path = thumb_base / photo_id / f"{size}.webp"
|
||||
if not thumb_path.exists():
|
||||
# Try user-prefixed paths: /data/thumbs/*/photo_id/size.webp
|
||||
matches = list(thumb_base.glob(f"*/{photo_id}/{size}.webp"))
|
||||
if matches:
|
||||
thumb_path = matches[0]
|
||||
@@ -79,7 +57,7 @@ def _load_thumb(photo_id: str, size: str = "medium") -> np.ndarray | None:
|
||||
return None
|
||||
try:
|
||||
img = Image.open(thumb_path).convert("RGB")
|
||||
img.load() # force decode to catch corruption early
|
||||
img.load()
|
||||
arr = np.array(img)
|
||||
img.close()
|
||||
return arr
|
||||
@@ -88,207 +66,25 @@ def _load_thumb(photo_id: str, size: str = "medium") -> np.ndarray | None:
|
||||
return None
|
||||
|
||||
|
||||
@shared_task(name='embed_photo', queue='vision', bind=True, max_retries=3)
|
||||
def embed_photo(self, photo_id: str):
|
||||
"""Generate CLIP embedding for a photo and store in pgvector."""
|
||||
if not is_enabled(FLAG_VISION_ENABLED):
|
||||
return {'status': 'skipped', 'reason': 'vision disabled'}
|
||||
|
||||
image = _load_thumb(photo_id, "medium") # 640px
|
||||
if image is None:
|
||||
return {'status': 'error', 'message': 'thumbnail not found'}
|
||||
|
||||
try:
|
||||
from app.services.vision.registry import registry
|
||||
embedder = registry.get_embedder()
|
||||
vector = embedder.embed_image(image)
|
||||
except Exception as exc:
|
||||
logger.exception("embed_photo failed for %s", photo_id)
|
||||
raise self.retry(exc=exc, countdown=60)
|
||||
|
||||
model_name = settings.vision.embedder.name
|
||||
|
||||
session = _get_sync_session()
|
||||
try:
|
||||
session.execute(
|
||||
delete(Embedding).where(
|
||||
Embedding.photo_id == photo_id,
|
||||
Embedding.model == model_name,
|
||||
)
|
||||
)
|
||||
emb = Embedding(
|
||||
photo_id=photo_id,
|
||||
model=model_name,
|
||||
vector=vector.tolist(),
|
||||
)
|
||||
session.add(emb)
|
||||
session.commit()
|
||||
except Exception:
|
||||
session.rollback()
|
||||
raise
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
logger.info("[%s] Embedded photo %s with %s", self.request.id, photo_id, model_name)
|
||||
return {'status': 'success', 'photo_id': photo_id}
|
||||
|
||||
|
||||
@shared_task(name='vision_fanout', queue='vision')
|
||||
def vision_fanout(photo_id: str):
|
||||
"""Dispatch all enabled vision tasks for a photo."""
|
||||
"""Dispatch vision work for a photo. Today this is just the binary
|
||||
classifier; the indirection stays so scanner/upload code keeps one
|
||||
entrypoint."""
|
||||
if not is_enabled(FLAG_VISION_ENABLED):
|
||||
return {'status': 'skipped', 'reason': 'vision disabled'}
|
||||
|
||||
embed_photo.delay(photo_id)
|
||||
|
||||
if is_enabled(FLAG_OCR_ENABLED):
|
||||
ocr_photo.delay(photo_id)
|
||||
if is_enabled(FLAG_DETECTOR_ENABLED):
|
||||
detect_objects.delay(photo_id)
|
||||
if is_enabled(FLAG_FACES_ENABLED):
|
||||
extract_faces.delay(photo_id)
|
||||
if is_enabled(FLAG_CLASSIFIER_ENABLED):
|
||||
classify_content.delay(photo_id)
|
||||
|
||||
classify_content.delay(photo_id)
|
||||
return {'status': 'dispatched', 'photo_id': photo_id}
|
||||
|
||||
|
||||
@shared_task(name='ocr_photo', queue='vision', bind=True, max_retries=3)
|
||||
def ocr_photo(self, photo_id: str):
|
||||
"""Run OCR on a photo and store text regions."""
|
||||
if not is_enabled(FLAG_VISION_ENABLED) or not is_enabled(FLAG_OCR_ENABLED):
|
||||
return {'status': 'skipped', 'reason': 'OCR disabled'}
|
||||
|
||||
image = _load_thumb(photo_id, "large") # 1280px for better OCR accuracy
|
||||
if image is None:
|
||||
return {'status': 'error', 'message': 'thumbnail not found'}
|
||||
|
||||
try:
|
||||
from app.services.vision.registry import registry
|
||||
ocr_engine = registry.get_ocr()
|
||||
results = ocr_engine.run(image)
|
||||
except Exception as exc:
|
||||
logger.exception("ocr_photo failed for %s", photo_id)
|
||||
raise self.retry(exc=exc, countdown=60)
|
||||
|
||||
if not results:
|
||||
logger.info("No OCR text found for photo %s", photo_id)
|
||||
return {'status': 'success', 'photo_id': photo_id, 'regions': 0}
|
||||
|
||||
from app.models.ocr_text import OCRText
|
||||
|
||||
session = _get_sync_session()
|
||||
try:
|
||||
session.execute(delete(OCRText).where(OCRText.photo_id == photo_id))
|
||||
for r in results:
|
||||
session.add(OCRText(
|
||||
photo_id=photo_id,
|
||||
text=r.text,
|
||||
language=r.language,
|
||||
confidence=r.confidence,
|
||||
bbox=r.bbox,
|
||||
))
|
||||
session.commit()
|
||||
except Exception:
|
||||
session.rollback()
|
||||
raise
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
logger.info("[%s] OCR: %d text regions for photo %s", self.request.id, len(results), photo_id)
|
||||
return {'status': 'success', 'photo_id': photo_id, 'regions': len(results)}
|
||||
|
||||
|
||||
@shared_task(name='detect_objects', queue='vision', bind=True, max_retries=3)
|
||||
def detect_objects(self, photo_id: str):
|
||||
"""Detect objects in a photo, create Tag(kind=object) rows, and
|
||||
link via photo_tags with confidence/bbox/source."""
|
||||
if not is_enabled(FLAG_VISION_ENABLED) or not is_enabled(FLAG_DETECTOR_ENABLED):
|
||||
return {'status': 'skipped', 'reason': 'detection disabled'}
|
||||
|
||||
image = _load_thumb(photo_id, "medium") # 640px
|
||||
if image is None:
|
||||
return {'status': 'error', 'message': 'thumbnail not found'}
|
||||
|
||||
try:
|
||||
from app.services.vision.registry import registry
|
||||
detector = registry.get_detector()
|
||||
detections = detector.detect(image)
|
||||
except Exception as exc:
|
||||
logger.exception("detect_objects failed for %s", photo_id)
|
||||
raise self.retry(exc=exc, countdown=60)
|
||||
|
||||
if not detections:
|
||||
logger.info("No objects detected for photo %s", photo_id)
|
||||
return {'status': 'success', 'photo_id': photo_id, 'objects': 0}
|
||||
|
||||
from app.models import Photo
|
||||
from app.models.tags import Tag, photo_tags
|
||||
|
||||
source_name = "vision:yolov8n"
|
||||
|
||||
session = _get_sync_session()
|
||||
try:
|
||||
# Get the photo's user_id so tags inherit ownership.
|
||||
photo = session.execute(
|
||||
select(Photo).where(Photo.id == photo_id)
|
||||
).scalar_one_or_none()
|
||||
owner_id = photo.user_id if photo else None
|
||||
|
||||
# Wipe previous detection results for this photo from this model
|
||||
session.execute(
|
||||
delete(photo_tags).where(
|
||||
photo_tags.c.photo_id == photo_id,
|
||||
photo_tags.c.source == source_name,
|
||||
)
|
||||
)
|
||||
|
||||
# Group detections by label, keep highest confidence per label
|
||||
best_per_label: dict[str, tuple[float, list]] = {}
|
||||
for det in detections:
|
||||
if det.label not in best_per_label or det.confidence > best_per_label[det.label][0]:
|
||||
best_per_label[det.label] = (det.confidence, det.bbox)
|
||||
|
||||
for label, (confidence, bbox) in best_per_label.items():
|
||||
# Find or create the object tag (scoped to user)
|
||||
tag = session.execute(
|
||||
select(Tag).where(Tag.name == label, Tag.kind == 'object', Tag.user_id == owner_id)
|
||||
).scalar_one_or_none()
|
||||
|
||||
if not tag:
|
||||
tag = Tag(name=label, kind='object', source=source_name, user_id=owner_id)
|
||||
session.add(tag)
|
||||
session.flush() # get tag.id
|
||||
|
||||
# Insert photo_tags association with ML metadata
|
||||
session.execute(
|
||||
photo_tags.insert().values(
|
||||
photo_id=photo_id,
|
||||
tag_id=tag.id,
|
||||
confidence=confidence,
|
||||
bbox=bbox,
|
||||
source=source_name,
|
||||
)
|
||||
)
|
||||
|
||||
session.commit()
|
||||
except Exception:
|
||||
session.rollback()
|
||||
raise
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
labels = [d.label for d in detections]
|
||||
logger.info("[%s] Detected %d objects in photo %s: %s", self.request.id, len(detections), photo_id, labels)
|
||||
return {'status': 'success', 'photo_id': photo_id, 'objects': len(detections)}
|
||||
|
||||
|
||||
@shared_task(name='classify_content', queue='vision', bind=True, max_retries=3)
|
||||
def classify_content(self, photo_id: str):
|
||||
"""Classify image content type (screenshot, document, artwork, etc.)
|
||||
using CLIP zero-shot classification. Writes Tag(kind=content_type)."""
|
||||
if not is_enabled(FLAG_VISION_ENABLED) or not is_enabled(FLAG_CLASSIFIER_ENABLED):
|
||||
return {'status': 'skipped', 'reason': 'classifier disabled'}
|
||||
"""Run the binary classifier and write:
|
||||
- a Tag(kind='content_type', name IN ('photography','other'))
|
||||
- Photo.needs_review = (label == 'other')
|
||||
"""
|
||||
if not is_enabled(FLAG_VISION_ENABLED):
|
||||
return {'status': 'skipped', 'reason': 'vision disabled'}
|
||||
|
||||
image = _load_thumb(photo_id, "medium")
|
||||
if image is None:
|
||||
@@ -297,30 +93,28 @@ def classify_content(self, photo_id: str):
|
||||
try:
|
||||
from app.services.vision.registry import registry
|
||||
classifier = registry.get_classifier()
|
||||
results = classifier.classify(image)
|
||||
result = classifier.classify(image)
|
||||
except Exception as exc:
|
||||
logger.exception("classify_content failed for %s", photo_id)
|
||||
raise self.retry(exc=exc, countdown=60)
|
||||
|
||||
if not results:
|
||||
logger.info("No confident classification for photo %s", photo_id)
|
||||
return {'status': 'success', 'photo_id': photo_id, 'content_type': None}
|
||||
|
||||
from app.models import Photo
|
||||
from app.models.tags import Tag, photo_tags
|
||||
|
||||
source_name = "vision:clip_classifier"
|
||||
best = results[0]
|
||||
label = result.label
|
||||
confidence = result.confidence
|
||||
|
||||
session = _get_sync_session()
|
||||
try:
|
||||
# Get the photo's user_id so tags inherit ownership.
|
||||
photo = session.execute(
|
||||
select(Photo).where(Photo.id == photo_id)
|
||||
).scalar_one_or_none()
|
||||
owner_id = photo.user_id if photo else None
|
||||
if photo is None:
|
||||
return {'status': 'error', 'message': 'photo not found'}
|
||||
owner_id = photo.user_id
|
||||
|
||||
# Wipe previous classification for this photo
|
||||
# Drop any previous classification for this photo.
|
||||
session.execute(
|
||||
delete(photo_tags).where(
|
||||
photo_tags.c.photo_id == photo_id,
|
||||
@@ -328,13 +122,13 @@ def classify_content(self, photo_id: str):
|
||||
)
|
||||
)
|
||||
|
||||
# Find or create content_type tag (scoped to user)
|
||||
tag = session.execute(
|
||||
select(Tag).where(Tag.name == best.label, Tag.kind == 'content_type', Tag.user_id == owner_id)
|
||||
select(Tag).where(
|
||||
Tag.name == label, Tag.kind == 'content_type', Tag.user_id == owner_id
|
||||
)
|
||||
).scalar_one_or_none()
|
||||
|
||||
if not tag:
|
||||
tag = Tag(name=best.label, kind='content_type', source=source_name, user_id=owner_id)
|
||||
tag = Tag(name=label, kind='content_type', source=source_name, user_id=owner_id)
|
||||
session.add(tag)
|
||||
session.flush()
|
||||
|
||||
@@ -342,232 +136,16 @@ def classify_content(self, photo_id: str):
|
||||
photo_tags.insert().values(
|
||||
photo_id=photo_id,
|
||||
tag_id=tag.id,
|
||||
confidence=best.confidence,
|
||||
confidence=confidence,
|
||||
source=source_name,
|
||||
)
|
||||
)
|
||||
session.commit()
|
||||
except Exception:
|
||||
session.rollback()
|
||||
raise
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
logger.info("[%s] Classified photo %s as '%s' (%.2f)", self.request.id, photo_id, best.label, best.confidence)
|
||||
return {'status': 'success', 'photo_id': photo_id, 'content_type': best.label}
|
||||
|
||||
|
||||
def _load_original(photo_id: str) -> np.ndarray | None:
|
||||
"""Load the original photo file as an RGB numpy array, resized to
|
||||
max 1280px on the longest edge for face detection."""
|
||||
from sqlalchemy import create_engine, select as sa_select, text as sa_text
|
||||
from app.models import Photo
|
||||
|
||||
session = _get_sync_session()
|
||||
try:
|
||||
photo = session.execute(
|
||||
sa_select(Photo).where(Photo.id == photo_id)
|
||||
).scalar_one_or_none()
|
||||
if not photo or not photo.filepath:
|
||||
return None
|
||||
filepath = photo.filepath
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
if not Path(filepath).exists():
|
||||
logger.warning("Original file not found: %s", filepath)
|
||||
return None
|
||||
|
||||
try:
|
||||
img = Image.open(filepath).convert("RGB")
|
||||
# Cap at 4000px on longest edge to avoid OOM, but keep as large
|
||||
# as possible for face detection accuracy
|
||||
max_dim = 4000
|
||||
w, h = img.size
|
||||
if max(w, h) > max_dim:
|
||||
scale = max_dim / max(w, h)
|
||||
resized = img.resize((int(w * scale), int(h * scale)), Image.BICUBIC)
|
||||
img.close()
|
||||
img = resized
|
||||
arr = np.array(img)
|
||||
img.close()
|
||||
return arr
|
||||
except Exception as e:
|
||||
logger.warning("Failed to load original %s: %s", filepath, e)
|
||||
return None
|
||||
|
||||
|
||||
@shared_task(name='extract_faces', queue='vision', bind=True, max_retries=3)
|
||||
def extract_faces(self, photo_id: str):
|
||||
"""Detect faces and store recognition embeddings using InsightFace
|
||||
(RetinaFace + ArcFace). No YOLO workaround needed — RetinaFace has
|
||||
strong human-vs-non-human precision on its own."""
|
||||
if not is_enabled(FLAG_VISION_ENABLED) or not is_enabled(FLAG_FACES_ENABLED):
|
||||
return {'status': 'skipped', 'reason': 'faces disabled'}
|
||||
|
||||
image = _load_original(photo_id)
|
||||
if image is None:
|
||||
image = _load_thumb(photo_id, "large")
|
||||
if image is None:
|
||||
return {'status': 'error', 'message': 'no image available'}
|
||||
|
||||
try:
|
||||
from app.services.vision.registry import registry
|
||||
face_proc = registry.get_face_processor()
|
||||
faces = face_proc.process(image)
|
||||
except Exception as exc:
|
||||
logger.exception("extract_faces failed for %s", photo_id)
|
||||
raise self.retry(exc=exc, countdown=60)
|
||||
|
||||
if not faces:
|
||||
logger.info("No faces detected for photo %s", photo_id)
|
||||
|
||||
return _save_faces(photo_id, faces)
|
||||
|
||||
|
||||
def _save_faces(photo_id: str, faces) -> dict:
|
||||
from app.models.face_embedding import FaceEmbedding
|
||||
|
||||
session = _get_sync_session()
|
||||
try:
|
||||
session.execute(delete(FaceEmbedding).where(FaceEmbedding.photo_id == photo_id))
|
||||
for face in faces:
|
||||
session.add(FaceEmbedding(
|
||||
photo_id=photo_id,
|
||||
bbox=face.bbox,
|
||||
vector=face.embedding.tolist(),
|
||||
quality=face.quality,
|
||||
cluster_id=None,
|
||||
))
|
||||
session.commit()
|
||||
except Exception:
|
||||
session.rollback()
|
||||
raise
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
if faces:
|
||||
logger.info("Extracted %d verified face(s) from photo %s", len(faces), photo_id)
|
||||
_schedule_recluster_debounced()
|
||||
return {'status': 'success', 'photo_id': photo_id, 'faces': len(faces)}
|
||||
|
||||
|
||||
RECLUSTER_DEBOUNCE_KEY = "mule:recluster_faces:pending"
|
||||
RECLUSTER_DELAY = 120 # seconds after last face extraction
|
||||
|
||||
|
||||
def _schedule_recluster_debounced():
|
||||
"""Schedule a recluster_faces run, debounced so rapid-fire face
|
||||
extractions don't spawn hundreds of redundant cluster jobs."""
|
||||
try:
|
||||
import redis as _redis
|
||||
r = _redis.from_url(settings.redis_url)
|
||||
already_pending = r.set(RECLUSTER_DEBOUNCE_KEY, "1",
|
||||
ex=RECLUSTER_DELAY, nx=True)
|
||||
if already_pending:
|
||||
recluster_faces.apply_async(countdown=RECLUSTER_DELAY)
|
||||
logger.info("Scheduled debounced recluster_faces in %ds", RECLUSTER_DELAY)
|
||||
except Exception as e:
|
||||
logger.debug("recluster debounce check failed: %s", e)
|
||||
|
||||
|
||||
@shared_task(name='recluster_faces', queue='vision', bind=True, max_retries=10)
|
||||
def recluster_faces(self):
|
||||
"""Run DBSCAN clustering over all face embeddings and assign/create
|
||||
Tag(kind=face_cluster) entries."""
|
||||
# Clear debounce key so new face extractions can schedule another round.
|
||||
try:
|
||||
import redis as _redis
|
||||
_redis.from_url(settings.redis_url).delete(RECLUSTER_DEBOUNCE_KEY)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if not _vision_worker_ready():
|
||||
logger.info("Vision worker not ready yet — retrying in 30s")
|
||||
raise self.retry(countdown=30)
|
||||
|
||||
if not is_enabled(FLAG_VISION_ENABLED) or not is_enabled(FLAG_FACES_ENABLED):
|
||||
return {'status': 'skipped', 'reason': 'faces disabled'}
|
||||
|
||||
from app.models import Photo
|
||||
from app.models.face_embedding import FaceEmbedding
|
||||
from app.models.tags import Tag, photo_tags
|
||||
from app.services.vision.clustering import cluster_faces
|
||||
|
||||
source_name = "vision:sface"
|
||||
|
||||
session = _get_sync_session()
|
||||
try:
|
||||
face_rows = session.execute(
|
||||
select(FaceEmbedding).order_by(FaceEmbedding.created_at)
|
||||
).scalars().all()
|
||||
|
||||
if len(face_rows) < 2:
|
||||
logger.info("Not enough faces for clustering (%d)", len(face_rows))
|
||||
return {'status': 'success', 'clusters': 0}
|
||||
|
||||
embeddings = np.array([f.vector for f in face_rows], dtype=np.float32)
|
||||
labels = cluster_faces(embeddings, eps=settings.vision.faces.cluster_eps)
|
||||
|
||||
# Clean up old face_cluster tags and their photo_tags
|
||||
old_cluster_tags = session.execute(
|
||||
select(Tag).where(Tag.kind == 'face_cluster', Tag.source == source_name)
|
||||
).scalars().all()
|
||||
for old_tag in old_cluster_tags:
|
||||
session.execute(
|
||||
delete(photo_tags).where(
|
||||
photo_tags.c.tag_id == old_tag.id,
|
||||
photo_tags.c.source == source_name,
|
||||
)
|
||||
)
|
||||
session.delete(old_tag)
|
||||
session.flush()
|
||||
|
||||
# Build new clusters
|
||||
cluster_tag_map: dict[int, str] = {}
|
||||
# Track which photos belong to which cluster
|
||||
cluster_photos: dict[int, set[str]] = {}
|
||||
|
||||
for i, label in enumerate(labels):
|
||||
if label == -1:
|
||||
face_rows[i].cluster_id = None
|
||||
continue
|
||||
|
||||
if label not in cluster_photos:
|
||||
cluster_photos[label] = set()
|
||||
cluster_photos[label].add(face_rows[i].photo_id)
|
||||
|
||||
if label not in cluster_tag_map:
|
||||
cluster_name = f"Person {label + 1}"
|
||||
# Inherit user_id from the representative photo.
|
||||
rep_photo = session.execute(
|
||||
select(Photo.user_id).where(Photo.id == face_rows[i].photo_id)
|
||||
).scalar_one_or_none()
|
||||
tag = Tag(
|
||||
name=cluster_name,
|
||||
kind='face_cluster',
|
||||
source=source_name,
|
||||
representative_photo_id=face_rows[i].photo_id,
|
||||
user_id=rep_photo,
|
||||
)
|
||||
session.add(tag)
|
||||
session.flush()
|
||||
cluster_tag_map[label] = tag.id
|
||||
|
||||
face_rows[i].cluster_id = cluster_tag_map[label]
|
||||
|
||||
# Write photo_tags associations so the tag count and tag_ids
|
||||
# filter work for face clusters
|
||||
for label, photo_ids in cluster_photos.items():
|
||||
tag_id = cluster_tag_map[label]
|
||||
for pid in photo_ids:
|
||||
session.execute(
|
||||
photo_tags.insert().values(
|
||||
photo_id=pid,
|
||||
tag_id=tag_id,
|
||||
source=source_name,
|
||||
)
|
||||
)
|
||||
session.execute(
|
||||
update(Photo)
|
||||
.where(Photo.id == photo_id)
|
||||
.values(needs_review=(label == 'other'))
|
||||
)
|
||||
|
||||
session.commit()
|
||||
except Exception:
|
||||
@@ -576,103 +154,41 @@ def recluster_faces(self):
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
n_clusters = len(cluster_tag_map)
|
||||
logger.info("Face clustering: %d clusters from %d faces", n_clusters, len(face_rows))
|
||||
return {'status': 'success', 'clusters': n_clusters, 'faces': len(face_rows)}
|
||||
logger.info("[%s] Classified %s as %s (%.2f)", self.request.id, photo_id, label, confidence)
|
||||
return {'status': 'success', 'photo_id': photo_id, 'label': label}
|
||||
|
||||
|
||||
@shared_task(name='backfill_vision', bind=True, max_retries=10)
|
||||
def backfill_vision(self, task: str | None = None, limit: int | None = None):
|
||||
"""Queue vision tasks for photos that haven't been processed yet.
|
||||
Uses a sync DB connection to avoid asyncpg conflicts in Celery."""
|
||||
def backfill_vision(self, limit: int | None = None, **_ignored):
|
||||
"""Queue classify_content for photos without a content_type tag."""
|
||||
if not _vision_worker_ready():
|
||||
logger.info("Vision worker not ready yet — retrying in 30s")
|
||||
raise self.retry(countdown=30)
|
||||
|
||||
model_name = settings.vision.embedder.name
|
||||
ordering = "ORDER BY p.taken_at DESC NULLS LAST, p.added_at DESC NULLS LAST"
|
||||
limit_clause = " LIMIT :lim" if limit else ""
|
||||
params: dict = {"model": model_name}
|
||||
params: dict = {}
|
||||
if limit:
|
||||
params["lim"] = int(limit)
|
||||
|
||||
session = _get_sync_session()
|
||||
try:
|
||||
# Each query finds photos missing a specific pipeline output so
|
||||
# enabling a new processor after import still back-fills.
|
||||
embed_ids = []
|
||||
if task in ('embed', None):
|
||||
sql = f"""
|
||||
SELECT p.id FROM photos p
|
||||
LEFT JOIN embeddings e ON e.photo_id = p.id AND e.model = :model
|
||||
WHERE e.photo_id IS NULL AND p.processing_status = 'completed'
|
||||
{ordering}{limit_clause}
|
||||
"""
|
||||
embed_ids = [r[0] for r in session.execute(sa_text(sql), params).fetchall()]
|
||||
|
||||
ocr_ids = []
|
||||
if task in ('ocr', None) and is_enabled(FLAG_OCR_ENABLED):
|
||||
sql = f"""
|
||||
SELECT p.id FROM photos p
|
||||
LEFT JOIN ocr_text o ON o.photo_id = p.id
|
||||
WHERE o.photo_id IS NULL AND p.processing_status = 'completed'
|
||||
{ordering}{limit_clause}
|
||||
"""
|
||||
ocr_ids = [r[0] for r in session.execute(sa_text(sql), params).fetchall()]
|
||||
|
||||
detect_ids = []
|
||||
if task in ('detect', None) and is_enabled(FLAG_DETECTOR_ENABLED):
|
||||
sql = f"""
|
||||
SELECT p.id FROM photos p
|
||||
WHERE p.processing_status = 'completed'
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM photo_tags pt WHERE pt.photo_id = p.id
|
||||
AND pt.source = 'vision:yolov8n'
|
||||
)
|
||||
{ordering}{limit_clause}
|
||||
"""
|
||||
detect_ids = [r[0] for r in session.execute(sa_text(sql), params).fetchall()]
|
||||
|
||||
face_ids = []
|
||||
if task in ('faces', None) and is_enabled(FLAG_FACES_ENABLED):
|
||||
sql = f"""
|
||||
SELECT p.id FROM photos p
|
||||
LEFT JOIN face_embeddings fe ON fe.photo_id = p.id
|
||||
WHERE fe.photo_id IS NULL AND p.processing_status = 'completed'
|
||||
{ordering}{limit_clause}
|
||||
"""
|
||||
face_ids = [r[0] for r in session.execute(sa_text(sql), params).fetchall()]
|
||||
|
||||
classify_ids = []
|
||||
if task in ('classify', None) and is_enabled(FLAG_CLASSIFIER_ENABLED):
|
||||
sql = f"""
|
||||
SELECT p.id FROM photos p
|
||||
WHERE p.processing_status = 'completed'
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM photo_tags pt WHERE pt.photo_id = p.id
|
||||
AND pt.source = 'vision:clip_classifier'
|
||||
)
|
||||
{ordering}{limit_clause}
|
||||
"""
|
||||
classify_ids = [r[0] for r in session.execute(sa_text(sql), params).fetchall()]
|
||||
except Exception:
|
||||
session.rollback()
|
||||
raise
|
||||
sql = f"""
|
||||
SELECT p.id FROM photos p
|
||||
WHERE p.processing_status = 'completed'
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM photo_tags pt
|
||||
WHERE pt.photo_id = p.id
|
||||
AND pt.source = 'vision:clip_classifier'
|
||||
)
|
||||
{ordering}{limit_clause}
|
||||
"""
|
||||
ids = [r[0] for r in session.execute(sa_text(sql), params).fetchall()]
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
# Dispatch — deduplicate across query results.
|
||||
all_ids = set(embed_ids) | set(ocr_ids) | set(detect_ids) | set(face_ids) | set(classify_ids)
|
||||
for pid in embed_ids:
|
||||
embed_photo.delay(pid)
|
||||
for pid in ocr_ids:
|
||||
ocr_photo.delay(pid)
|
||||
for pid in detect_ids:
|
||||
detect_objects.delay(pid)
|
||||
for pid in face_ids:
|
||||
extract_faces.delay(pid)
|
||||
for pid in classify_ids:
|
||||
for pid in ids:
|
||||
classify_content.delay(pid)
|
||||
|
||||
logger.info("Backfill queued %d photos for vision processing", len(all_ids))
|
||||
return {'status': 'queued', 'count': len(all_ids)}
|
||||
logger.info("Backfill queued %d photos for classification", len(ids))
|
||||
return {'status': 'queued', 'count': len(ids)}
|
||||
|
||||
Reference in New Issue
Block a user