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

"""
Celery tasks for the vision pipeline.

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

from app.config import settings
from app.services.feature_flags import is_enabled, FLAG_VISION_ENABLED

logger = logging.getLogger(__name__)

VISION_READY_KEY = "mulita:vision:ready"


def _vision_worker_ready() -> bool:
    try:
        import redis as _redis
        return bool(_redis.from_url(settings.redis_url).exists(VISION_READY_KEY))
    except Exception:
        return False


_sync_engine = None


def _get_sync_engine():
    global _sync_engine
    if _sync_engine is None:
        sync_url = settings.database_url.replace("+asyncpg", "+psycopg2").replace("+aiosqlite", "")
        _sync_engine = create_engine(sync_url, pool_pre_ping=True, pool_size=3, max_overflow=5)
    return _sync_engine


def _get_sync_session() -> Session:
    return sessionmaker(bind=_get_sync_engine())()


def _load_thumb(photo_id: str, size: str = "medium") -> np.ndarray | None:
    thumb_base = Path("/data/thumbs")
    thumb_path = thumb_base / photo_id / f"{size}.webp"
    if not thumb_path.exists():
        matches = list(thumb_base.glob(f"*/{photo_id}/{size}.webp"))
        if matches:
            thumb_path = matches[0]
        else:
            logger.warning("Thumbnail not found: %s", thumb_path)
            return None
    try:
        img = Image.open(thumb_path).convert("RGB")
        img.load()
        arr = np.array(img)
        img.close()
        return arr
    except Exception as e:
        logger.warning("Corrupt or unreadable thumbnail for %s: %s", photo_id, e)
        return None


@shared_task(name='vision_fanout', queue='vision')
def vision_fanout(photo_id: str):
    """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'}
    classify_content.delay(photo_id)
    return {'status': 'dispatched', 'photo_id': photo_id}


@shared_task(name='classify_content', queue='vision', bind=True, max_retries=3)
def classify_content(self, photo_id: str):
    """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:
        return {'status': 'error', 'message': 'thumbnail not found'}

    try:
        from app.services.vision.registry import registry
        classifier = registry.get_classifier()
        result = classifier.classify(image)
    except Exception as exc:
        logger.exception("classify_content failed for %s", photo_id)
        raise self.retry(exc=exc, countdown=60)

    from app.models import Photo
    from app.models.tags import Tag, photo_tags

    source_name = "vision:clip_classifier"
    label = result.label
    confidence = result.confidence

    session = _get_sync_session()
    try:
        photo = session.execute(
            select(Photo).where(Photo.id == photo_id)
        ).scalar_one_or_none()
        if photo is None:
            return {'status': 'error', 'message': 'photo not found'}
        owner_id = photo.user_id

        # Drop any previous classification for this photo.
        session.execute(
            delete(photo_tags).where(
                photo_tags.c.photo_id == photo_id,
                photo_tags.c.source == source_name,
            )
        )

        tag = session.execute(
            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=label, kind='content_type', source=source_name, user_id=owner_id)
            session.add(tag)
            session.flush()

        session.execute(
            photo_tags.insert().values(
                photo_id=photo_id,
                tag_id=tag.id,
                confidence=confidence,
                source=source_name,
            )
        )

        session.execute(
            update(Photo)
            .where(Photo.id == photo_id)
            .values(needs_review=(label == 'other'))
        )

        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()

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

    ordering = "ORDER BY p.taken_at DESC NULLS LAST, p.added_at DESC NULLS LAST"
    limit_clause = " LIMIT :lim" if limit else ""
    params: dict = {}
    if limit:
        params["lim"] = int(limit)

    session = _get_sync_session()
    try:
        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()

    for pid in ids:
        classify_content.delay(pid)

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