app.py

"""
VERITAS Backend — AI Image Detection Engine v2
Ensemble: 2 models + EXIF metadata check + preprocessing
"""

import io
import logging
import struct

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from PIL import Image, ExifTags
from transformers import pipeline

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("veritas")

MODELS = [
    "umm-maybe/AI-image-detector",
    "prithivMLmods/Deep-Fake-Detector-v2-Model",
]
MODEL_WEIGHTS = [0.45, 0.55]   # weight toward newer model
MAX_FILE_SIZE_MB = 25
ALLOWED_CONTENT_TYPES = {"image/jpeg", "image/png", "image/webp", "image/bmp"}
INPUT_SIZE = 224   # normalize input to both models

app = FastAPI(title="VERITAS Detection Engine v2")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

detectors = []


@app.on_event("startup")
def load_models():
    global detectors
    for m in MODELS:
        logger.info(f"Loading: {m}")
        detectors.append(pipeline("image-classification", model=m))
    logger.info("All models loaded.")


@app.get("/health")
def health():
    return {"status": "ok", "models_loaded": len(detectors)}


def get_ai_score(results: list) -> float:
    """Extract AI/fake probability from classifier output."""
    score = next(
        (r["score"] for r in results
         if any(k in r["label"].lower() for k in ("fake", "ai", "artificial", "generated", "deepfake"))),
        None
    )
    if score is None:
        score = next(
            (r["score"] for r in results
             if not any(k in r["label"].lower() for k in ("real", "human", "authentic"))),
            results[0]["score"]
        )
    return score


def exif_penalty(image_bytes: bytes) -> float:
    """
    Return a small upward nudge (0–8%) if image lacks real camera EXIF.
    AI-generated images almost never have camera maker/model tags.
    """
    try:
        img = Image.open(io.BytesIO(image_bytes))
        exif_data = img._getexif()
        if not exif_data:
            return 5.0   # no EXIF at all — slight nudge up
        tags = {ExifTags.TAGS.get(k, k): v for k, v in exif_data.items()}
        has_camera = "Make" in tags or "Model" in tags
        has_datetime = "DateTimeOriginal" in tags
        if not has_camera and not has_datetime:
            return 4.0
        if not has_camera:
            return 2.0
        return 0.0       # legit camera EXIF present — no nudge
    except Exception:
        return 3.0       # can't read EXIF — small nudge


@app.post("/detect")
async def detect(file: UploadFile = File(...)):
    if not detectors:
        raise HTTPException(status_code=503, detail="Models not loaded yet")

    if file.content_type not in ALLOWED_CONTENT_TYPES:
        raise HTTPException(status_code=415, detail=f"Unsupported type: {file.content_type}")

    raw = await file.read()
    if len(raw) > MAX_FILE_SIZE_MB * 1024 * 1024:
        raise HTTPException(status_code=413, detail=f"File exceeds {MAX_FILE_SIZE_MB}MB")

    try:
        image = Image.open(io.BytesIO(raw)).convert("RGB")
        # Step 3: normalize input size for consistent inference
        image = image.resize((INPUT_SIZE, INPUT_SIZE), Image.LANCZOS)
    except Exception:
        raise HTTPException(status_code=400, detail="Could not decode image")

    # Step 2: run ensemble
    raw_scores = []
    all_results = []
    try:
        for det in detectors:
            res = det(image)
            all_results.append(res)
            raw_scores.append(get_ai_score(res))
    except Exception as e:
        logger.exception("Inference failed")
        raise HTTPException(status_code=500, detail=f"Inference error: {e}")

    # Weighted average of model scores
    ensemble_score = sum(s * w for s, w in zip(raw_scores, MODEL_WEIGHTS))

    # Step 4: EXIF penalty (additive, capped so total ≤ 100)
    penalty = exif_penalty(raw)
    final_score = min(ensemble_score * 100 + penalty, 100.0)
    percentage = round(final_score, 1)

    return JSONResponse({
        "ai_probability": percentage,
        "model_scores": [round(s * 100, 1) for s in raw_scores],
        "exif_penalty": penalty,
        "filename": file.filename,
    })