app/detector.py

"""
Detector API routes.
"""

import logging
from functools import lru_cache
from io import BytesIO

from typing import Callable

from fastapi import APIRouter, Depends, File, HTTPException, UploadFile
from PIL import Image, UnidentifiedImageError

from app.core.detector.model import (
    SiglipResources,
    ViTResources,
    get_siglip_model,
    get_vit_model,
)
from app.services.detector.prediction import predict_single_image
from app.services.detector.transforms import get_eval_transforms

detector = APIRouter()
logger = logging.getLogger(__name__)


@lru_cache(maxsize=1)
def get_siglip_transforms():
    """Build and cache SigLIP evaluation transforms once per process."""

    resources = get_siglip_model()
    return get_eval_transforms(resources.processor, "siglip")


@lru_cache(maxsize=1)
def get_vit_transforms():
    """Build and cache ViT evaluation transforms once per process."""

    resources = get_vit_model()
    return get_eval_transforms(resources.processor, "vit")


@detector.post("/siglip-detect")
async def siglip_detect_deepfake(
    file: UploadFile = File(...),
    resources: SiglipResources = Depends(get_siglip_model),
    siglip_transforms: Callable = Depends(get_siglip_transforms),
):
    """
    Detect if an image is a deepfake or real using SigLIP + LoRA model.

    Args:
        file: Uploaded image file

    Returns:
        JSON response with prediction results
    """

    try:
        image_bytes = await file.read()
        image = Image.open(BytesIO(image_bytes)).convert("RGB")

        result = predict_single_image(
            image, resources.model, siglip_transforms, "SigLIP + LoRA"
        )

        return result
    except UnidentifiedImageError:
        raise HTTPException(status_code=422, detail="Invalid or unsupported image file")
    except HTTPException:
        raise
    except Exception as exc:  # pragma: no cover - defensive server guard
        logger.exception("Unhandled error during deepfake detection")
        raise HTTPException(status_code=500, detail="Error processing image") from exc


@detector.post("/vit-detect")
async def vit_detect_deepfake(
    file: UploadFile = File(...),
    resources: ViTResources = Depends(get_vit_model),
    vit_transforms: Callable = Depends(get_vit_transforms),
):
    """
    Detect if an image is a deepfake or real using ViT + LoRA model.

    Args:
        file: Uploaded image file

    Returns:
        JSON response with prediction results
    """

    try:
        image_bytes = await file.read()
        image = Image.open(BytesIO(image_bytes)).convert("RGB")

        result = predict_single_image(
            image, resources.model, vit_transforms, "ViT + LoRA"
        )

        return result
    except UnidentifiedImageError:
        raise HTTPException(status_code=422, detail="Invalid or unsupported image file")
    except HTTPException:
        raise
    except Exception as exc:  # pragma: no cover - defensive server guard
        logger.exception("Unhandled error during deepfake detection")
        raise HTTPException(status_code=500, detail="Error processing image") from exc