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)}


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


def _iou(a: list[float], b: list[float]) -> float:
    """Intersection-over-area of box a within box b (how much of a is inside b).
    Boxes are [x1, y1, x2, y2] normalized 0-1."""
    x1 = max(a[0], b[0])
    y1 = max(a[1], b[1])
    x2 = min(a[2], b[2])
    y2 = min(a[3], b[3])
    inter = max(0, x2 - x1) * max(0, y2 - y1)
    area_a = max(0, a[2] - a[0]) * max(0, a[3] - a[1])
    return inter / area_a if area_a > 0 else 0


def _face_inside_person(face_bbox: list[float], person_bboxes: list[list[float]]) -> bool:
    """Return True if the face bbox overlaps at least 50% with any YOLO
    'person' detection. Filters out faces on dogs, paintings, etc."""
    for pb in person_bboxes:
        if _iou(face_bbox, pb) >= 0.5:
            return True
    return False


@shared_task(name='extract_faces', queue='vision')
def extract_faces(photo_id: str):
    """Detect faces and store recognition embeddings. Only keeps faces
    that overlap with a YOLO 'person' detection to filter out animal
    and cartoon false positives."""
    if not settings.vision.enabled or not settings.vision.faces.enabled:
        return {'status': 'skipped', 'reason': 'faces disabled'}

    # Use original file for face detection — thumbnails are often too
    # small (240px) for reliable face detection.
    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'}

    # Step 1: run YOLO to find "person" bounding boxes
    from app.services.vision.registry import registry
    detector = registry.get_detector()
    thumb = _load_thumb(photo_id, "medium")
    person_bboxes = []
    if thumb is not None:
        detections = detector.detect(thumb)
        person_bboxes = [d.bbox for d in detections if d.label == 'person']

    # Step 2: run face detection
    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, [])

    # Step 3: filter — keep only faces inside a person bbox
    if person_bboxes:
        verified = [f for f in faces if _face_inside_person(f.bbox, person_bboxes)]
        dropped = len(faces) - len(verified)
        if dropped:
            logger.info("Dropped %d non-person face(s) for photo %s", dropped, photo_id)
        faces = verified
    else:
        # No person detected by YOLO → drop all face detections
        # (no human body visible = likely false positives)
        logger.info("No YOLO person detected, dropping %d face(s) for photo %s", len(faces), photo_id)
        faces = []

    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
    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}