mule-image/backend/app/tasks/vision.py

"""
Celery tasks for the vision pipeline — embedding, OCR, object detection,
face recognition.

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.
"""
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.orm import Session, sessionmaker
from PIL import Image

from app.models.embeddings import Embedding
from app.config import settings

logger = logging.getLogger(__name__)


def _get_sync_session() -> Session:
    """Create a sync DB session for use in Celery workers."""
    sync_url = settings.database_url.replace("+asyncpg", "+psycopg2").replace("+aiosqlite", "")
    engine = create_engine(sync_url, pool_pre_ping=True)
    return sessionmaker(bind=engine)()


def _load_thumb(photo_id: str, size: str = "medium") -> np.ndarray | None:
    """Load a thumbnail as an RGB numpy array."""
    thumb_path = Path(f"/data/thumbs/{photo_id}/{size}.webp")
    if not thumb_path.exists():
        logger.warning("Thumbnail not found: %s", thumb_path)
        return None
    img = Image.open(thumb_path).convert("RGB")
    return np.array(img)


@shared_task(name='embed_photo', queue='vision')
def embed_photo(photo_id: str):
    """Generate CLIP embedding for a photo and store in pgvector."""
    if not settings.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'}

    from app.services.vision.registry import registry
    embedder = registry.get_embedder()
    vector = embedder.embed_image(image)

    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()
    finally:
        session.close()

    logger.info("Embedded photo %s with %s", 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."""
    if not settings.vision.enabled:
        return {'status': 'skipped', 'reason': 'vision disabled'}

    embed_photo.delay(photo_id)

    if settings.vision.ocr.enabled:
        ocr_photo.delay(photo_id)
    if settings.vision.detector.enabled:
        detect_objects.delay(photo_id)
    if settings.vision.faces.enabled:
        extract_faces.delay(photo_id)

    return {'status': 'dispatched', 'photo_id': photo_id}


@shared_task(name='ocr_photo', queue='vision')
def ocr_photo(photo_id: str):
    """Run OCR on a photo and store text regions."""
    if not settings.vision.enabled or not settings.vision.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'}

    from app.services.vision.registry import registry
    ocr_engine = registry.get_ocr()
    results = ocr_engine.run(image)

    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()
    finally:
        session.close()

    logger.info("OCR: %d text regions for photo %s", len(results), photo_id)
    return {'status': 'success', 'photo_id': photo_id, 'regions': len(results)}


@shared_task(name='detect_objects', queue='vision')
def detect_objects(photo_id: str):
    """Detect objects in a photo, create Tag(kind=object) rows, and
    link via photo_tags with confidence/bbox/source."""
    if not settings.vision.enabled or not settings.vision.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'}

    from app.services.vision.registry import registry
    detector = registry.get_detector()
    detections = detector.detect(image)

    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.tags import Tag, photo_tags

    source_name = "vision:yolov8n"

    session = _get_sync_session()
    try:
        # 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
            tag = session.execute(
                select(Tag).where(Tag.name == label, Tag.kind == 'object')
            ).scalar_one_or_none()

            if not tag:
                tag = Tag(name=label, kind='object', source=source_name)
                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()
    finally:
        session.close()

    labels = [d.label for d in detections]
    logger.info("Detected %d objects in photo %s: %s", len(detections), photo_id, labels)
    return {'status': 'success', 'photo_id': photo_id, 'objects': len(detections)}


@shared_task(name='extract_faces', queue='vision')
def extract_faces(photo_id: str):
    """Detect faces and store recognition embeddings. Clustering is
    handled separately by recluster_faces (periodic task)."""
    if not settings.vision.enabled or not settings.vision.faces.enabled:
        return {'status': 'skipped', 'reason': 'faces disabled'}

    image = _load_thumb(photo_id, "large")  # 1280px for better face detection
    if image is None:
        return {'status': 'error', 'message': 'thumbnail not found'}

    from app.services.vision.registry import registry
    face_proc = registry.get_face_processor()
    faces = face_proc.process(image)

    if not faces:
        logger.info("No faces detected for photo %s", photo_id)
        return {'status': 'success', 'photo_id': photo_id, 'faces': 0}

    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()
    finally:
        session.close()

    logger.info("Extracted %d faces from photo %s", len(faces), photo_id)
    return {'status': 'success', 'photo_id': photo_id, 'faces': len(faces)}


@shared_task(name='recluster_faces', queue='vision')
def recluster_faces():
    """Run DBSCAN clustering over all face embeddings and assign/create
    Tag(kind=face_cluster) entries."""
    if not settings.vision.enabled or not settings.vision.faces.enabled:
        return {'status': 'skipped', 'reason': 'faces disabled'}

    from app.models.face_embedding import FaceEmbedding
    from app.models.tags import Tag
    from app.services.vision.clustering import cluster_faces

    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)

        cluster_tag_map: dict[int, str] = {}
        source_name = "vision:sface"

        for i, label in enumerate(labels):
            if label == -1:
                face_rows[i].cluster_id = None
                continue

            if label not in cluster_tag_map:
                cluster_name = f"Person {label + 1}"
                tag = session.execute(
                    select(Tag).where(
                        Tag.kind == 'face_cluster',
                        Tag.source == source_name,
                        Tag.name == cluster_name,
                    )
                ).scalar_one_or_none()

                if not tag:
                    tag = Tag(
                        name=cluster_name,
                        kind='face_cluster',
                        source=source_name,
                        representative_photo_id=face_rows[i].photo_id,
                    )
                    session.add(tag)
                    session.flush()
                cluster_tag_map[label] = tag.id

            face_rows[i].cluster_id = cluster_tag_map[label]

        session.commit()
    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)}


@shared_task(name='backfill_vision')
def backfill_vision(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."""
    model_name = settings.vision.embedder.name
    sql = """
        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'
        ORDER BY p.added_at DESC
    """
    if limit:
        sql += f" LIMIT {limit}"

    session = _get_sync_session()
    try:
        result = session.execute(sa_text(sql), {"model": model_name})
        photo_ids = [row[0] for row in result.fetchall()]
    finally:
        session.close()

    count = 0
    for pid in photo_ids:
        if task == 'embed' or task is None:
            embed_photo.delay(pid)
        if task == 'ocr' or task is None:
            ocr_photo.delay(pid)
        if task == 'detect' or task is None:
            detect_objects.delay(pid)
        if task == 'faces' or task is None:
            extract_faces.delay(pid)
        count += 1

    logger.info("Backfill queued %d photos for vision processing", count)
    return {'status': 'queued', 'count': count}