app/main.py

import io
import tempfile
from pathlib import Path

from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.staticfiles import StaticFiles
from PIL import Image

from .model import load_detector, predict_image
from .video import sample_frames

MAX_IMAGE_SIZE_MB = 20
MAX_VIDEO_SIZE_MB = 100
N_VIDEO_FRAMES = 5

IMAGE_TYPES = {"image/jpeg", "image/jpg", "image/png", "image/webp"}
VIDEO_TYPES = {"video/mp4", "video/quicktime", "video/webm", "video/x-matroska"}

app = FastAPI(title="Deepfake Detector")


@app.on_event("startup")
def warmup():
    load_detector()


@app.post("/api/predict")
async def predict(file: UploadFile = File(...)):
    content_type = (file.content_type or "").lower()
    raw = await file.read()
    size_mb = len(raw) / (1024 * 1024)

    if content_type in IMAGE_TYPES:
        if size_mb > MAX_IMAGE_SIZE_MB:
            raise HTTPException(413, f"Image exceeds {MAX_IMAGE_SIZE_MB} MB")
        try:
            image = Image.open(io.BytesIO(raw))
        except Exception:
            raise HTTPException(400, "Invalid image")
        p_fake = predict_image(image)
        return {
            "media_type": "image",
            "p_fake": p_fake,
            "reliability": 1.0 - p_fake,
            "n_frames": 1,
        }

    if content_type in VIDEO_TYPES:
        if size_mb > MAX_VIDEO_SIZE_MB:
            raise HTTPException(413, f"Video exceeds {MAX_VIDEO_SIZE_MB} MB")
        suffix = Path(file.filename or "video.mp4").suffix or ".mp4"
        with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
            tmp.write(raw)
            tmp_path = tmp.name
        try:
            frames = sample_frames(tmp_path, N_VIDEO_FRAMES)
        except ValueError as e:
            raise HTTPException(400, str(e))
        finally:
            try:
                Path(tmp_path).unlink(missing_ok=True)
            except Exception:
                pass
        probs = [predict_image(f) for f in frames]
        p_fake = sum(probs) / len(probs)
        return {
            "media_type": "video",
            "p_fake": p_fake,
            "reliability": 1.0 - p_fake,
            "n_frames": len(frames),
            "frame_probs": probs,
        }

    raise HTTPException(415, f"Unsupported media type: {content_type}")


static_dir = Path(__file__).parent / "static"
app.mount("/", StaticFiles(directory=str(static_dir), html=True), name="static")