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)
    if settings.vision.classifier.enabled:
        classify_content.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='classify_content', queue='vision')
def classify_content(photo_id: str):
    """Classify image content type (screenshot, document, artwork, etc.)
    using CLIP zero-shot classification. Writes Tag(kind=content_type)."""
    if not settings.vision.enabled or not settings.vision.classifier.enabled:
        return {'status': 'skipped', 'reason': 'classifier disabled'}

    image = _load_thumb(photo_id, "medium")
    if image is None:
        return {'status': 'error', 'message': 'thumbnail not found'}

    from app.services.vision.registry import registry
    classifier = registry.get_classifier()
    results = classifier.classify(image)

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

    source_name = "vision:clip_classifier"
    best = results[0]

    session = _get_sync_session()
    try:
        # Wipe previous classification for this photo
        session.execute(
            delete(photo_tags).where(
                photo_tags.c.photo_id == photo_id,
                photo_tags.c.source == source_name,
            )
        )

        # Find or create content_type tag
        tag = session.execute(
            select(Tag).where(Tag.name == best.label, Tag.kind == 'content_type')
        ).scalar_one_or_none()

        if not tag:
            tag = Tag(name=best.label, kind='content_type', source=source_name)
            session.add(tag)
            session.flush()

        session.execute(
            photo_tags.insert().values(
                photo_id=photo_id,
                tag_id=tag.id,
                confidence=best.confidence,
                source=source_name,
            )
        )
        session.commit()
    finally:
        session.close()

    logger.info("Classified photo %s as '%s' (%.2f)", 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)
            img = img.resize((int(w * scale), int(h * scale)), Image.BICUBIC)
        return np.array(img)
    except Exception as e:
        logger.warning("Failed to load original %s: %s", filepath, e)
        return None


@shared_task(name='extract_faces', queue='vision')
def extract_faces(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 settings.vision.enabled or not settings.vision.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'}

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

    if faces:
        logger.info("Extracted %d verified face(s) 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, 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}"
                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]

        # 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.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
    # Newest-first ordering — matches regenerate_all_thumbnails so the
    # whole ingestion pipeline sweeps the library top-down and the user
    # sees recent photos fully-indexed long before the backlog drains.
    # `taken_at` is the canonical capture timestamp (from EXIF, falls
    # back to filesystem mtime in scan); `added_at` is the tie-breaker
    # when taken_at is null.
    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.taken_at DESC NULLS LAST, p.added_at DESC NULLS LAST
    """
    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}