video_editing_poc/tools/local_cv_worker.py

#!/usr/bin/env python3
"""Optional local CV worker for the Spring visual-analysis provider.

Run:
  pip install -r tools/local_cv_requirements.txt
  uvicorn tools.local_cv_worker:app --host 127.0.0.1 --port 8091
"""

from __future__ import annotations

import os
import logging
import time
from pathlib import Path
from typing import Any

try:
    import cv2  # type: ignore
except ImportError:  # pragma: no cover - optional runtime dependency
    cv2 = None

try:
    from fastapi import FastAPI
except ImportError as exc:  # pragma: no cover - fail clearly at worker startup
    raise SystemExit("Install FastAPI first: pip install fastapi uvicorn") from exc

try:
    from ultralytics import YOLO  # type: ignore
except ImportError:  # pragma: no cover - optional runtime dependency
    YOLO = None


app = FastAPI(title="Local CV Visual Analysis Worker")
_yolo_model: Any | None = None
logging.basicConfig(level=os.getenv("LOCAL_CV_LOG_LEVEL", "INFO"))
log = logging.getLogger("local_cv_worker")


@app.get("/health")
def health() -> dict[str, Any]:
    return {
        "status": "ok",
        "opencvAvailable": cv2 is not None,
        "yoloAvailable": YOLO is not None,
        "yoloModel": configured_yolo_model(),
    }


@app.post("/v1/analyze-visuals")
def analyze_visuals(payload: dict[str, Any]) -> dict[str, Any]:
    started_at = time.monotonic()
    source = payload.get("source", {})
    thumbnails = [str(item) for item in payload.get("thumbnails", [])]
    readable_thumbnails = [path for path in thumbnails if Path(path).is_file()]
    clip_id = source.get("clipId", "source")
    log.info(
        "event=local_cv_worker_request_started clip_id=%s thumbnails=%d readable_thumbnails=%d "
        "shot_segments=%d duration_seconds=%s",
        clip_id,
        len(thumbnails),
        len(readable_thumbnails),
        len(payload.get("shotSegments", [])),
        source.get("durationSeconds", 0),
    )
    if len(readable_thumbnails) != len(thumbnails):
        unreadable = [path for path in thumbnails if path not in readable_thumbnails]
        log.warning(
            "event=local_cv_worker_unreadable_thumbnails clip_id=%s unreadable_count=%d unreadable=%s",
            clip_id,
            len(unreadable),
            unreadable[:10],
        )
    quality = quality_scores(readable_thumbnails)
    labels = object_labels(readable_thumbnails)

    response = {
        "clipId": clip_id,
        "blurScore": quality["blurScore"],
        "exposureScore": quality["exposureScore"],
        "motionScore": motion_score(payload.get("shotSegments", []), source.get("durationSeconds", 0)),
        "compositionScore": composition_score(source, readable_thumbnails),
        "facePresence": face_presence(readable_thumbnails),
        "objectLabels": labels or [{"label": "unknown", "confidence": 0.2, "source": "local_cv_worker"}],
        "representativeThumbnails": representative_thumbnails(readable_thumbnails or thumbnails),
        "analysisMethod": "local_cv_worker_opencv_yolo" if labels else "local_cv_worker_opencv",
    }
    log.info(
        "event=local_cv_worker_request_completed clip_id=%s elapsed_ms=%d blur_score=%s "
        "exposure_score=%s object_labels=%d method=%s",
        clip_id,
        elapsed_ms(started_at),
        response["blurScore"],
        response["exposureScore"],
        len(response["objectLabels"]),
        response["analysisMethod"],
    )
    return response


def quality_scores(thumbnails: list[str]) -> dict[str, float]:
    if cv2 is None or not thumbnails:
        log.info(
            "event=local_cv_worker_quality_defaulted reason=%s",
            "opencv_unavailable" if cv2 is None else "no_readable_thumbnails",
        )
        return {"blurScore": 0.5, "exposureScore": 0.5}

    blur_scores: list[float] = []
    exposure_scores: list[float] = []
    for thumbnail in thumbnails:
        image = cv2.imread(thumbnail)
        if image is None:
            continue
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        laplacian_variance = float(cv2.Laplacian(gray, cv2.CV_64F).var())
        blur_scores.append(clamp(laplacian_variance / 500.0))
        mean_brightness = float(gray.mean()) / 255.0
        exposure_scores.append(clamp(1.0 - abs(mean_brightness - 0.5) * 2.0))

    return {
        "blurScore": average_or_default(blur_scores, 0.5),
        "exposureScore": average_or_default(exposure_scores, 0.5),
    }


def object_labels(thumbnails: list[str]) -> list[dict[str, Any]]:
    started_at = time.monotonic()
    model = yolo_model()
    if model is None or not thumbnails:
        log.info(
            "event=local_cv_worker_object_detection_skipped reason=%s thumbnails=%d",
            "model_unavailable" if model is None else "no_readable_thumbnails",
            len(thumbnails),
        )
        return []

    labels: dict[str, float] = {}
    for result in model(thumbnails, verbose=False):
        names = getattr(result, "names", {})
        boxes = getattr(result, "boxes", None)
        if boxes is None:
            continue
        for box in boxes:
            label = str(names.get(int(box.cls[0]), "object"))
            confidence = float(box.conf[0])
            labels[label] = max(labels.get(label, 0.0), confidence)

    detected = [
        {"label": label, "confidence": round(clamp(confidence), 3), "source": "yolo"}
        for label, confidence in sorted(labels.items(), key=lambda item: item[1], reverse=True)[:10]
    ]
    log.info(
        "event=local_cv_worker_object_detection_completed elapsed_ms=%d thumbnails=%d labels=%d",
        elapsed_ms(started_at),
        len(thumbnails),
        len(detected),
    )
    return detected


def yolo_model() -> Any | None:
    global _yolo_model
    model_path = configured_yolo_model()
    if YOLO is None or not model_path:
        return None
    if _yolo_model is None:
        started_at = time.monotonic()
        log.info("event=local_cv_worker_yolo_model_loading model=%s", model_path)
        _yolo_model = YOLO(model_path)
        log.info("event=local_cv_worker_yolo_model_loaded model=%s elapsed_ms=%d", model_path, elapsed_ms(started_at))
    return _yolo_model


def configured_yolo_model() -> str:
    if os.getenv("LOCAL_CV_DISABLE_YOLO", "false").lower() == "true":
        return ""
    return os.getenv("LOCAL_CV_YOLO_MODEL", "yolov8n.pt")


def face_presence(thumbnails: list[str]) -> str:
    if cv2 is None or not thumbnails:
        return "unknown_without_face_detector"
    cascade_path = getattr(cv2.data, "haarcascades", "") + "haarcascade_frontalface_default.xml"
    if not Path(cascade_path).is_file():
        return "unknown_without_face_detector"
    cascade = cv2.CascadeClassifier(cascade_path)
    for thumbnail in thumbnails:
        image = cv2.imread(thumbnail)
        if image is None:
            continue
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        faces = cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=4)
        if len(faces) > 0:
            return "faces_detected"
    return "no_faces_detected"


def motion_score(shot_segments: list[dict[str, Any]], duration_seconds: float) -> float:
    try:
        duration = float(duration_seconds)
    except (TypeError, ValueError):
        duration = 0.0
    if duration <= 0:
        return 0.3
    cuts_per_minute = max(0, len(shot_segments) - 1) / duration * 60.0
    return round(clamp(0.25 + cuts_per_minute / 20.0), 3)


def composition_score(source: dict[str, Any], thumbnails: list[str]) -> float:
    width = float(source.get("width", 0) or 0)
    height = float(source.get("height", 0) or 0)
    aspect_ratio = width / height if height else 0
    aspect_score = 0.8 if 1.70 <= aspect_ratio <= 1.90 else 0.55
    thumbnail_score = 0.7 if thumbnails else 0.45
    return round((aspect_score + thumbnail_score) / 2.0, 3)


def representative_thumbnails(thumbnails: list[str]) -> list[str]:
    if len(thumbnails) <= 3:
        return thumbnails
    return [thumbnails[0], thumbnails[len(thumbnails) // 2], thumbnails[-1]]


def average_or_default(values: list[float], default: float) -> float:
    if not values:
        return default
    return round(sum(values) / len(values), 3)


def clamp(value: float) -> float:
    return max(0.0, min(1.0, value))


def elapsed_ms(started_at: float) -> int:
    return int((time.monotonic() - started_at) * 1000)