from fastapi import FastAPI, Request
import requests
from PIL import Image
from io import BytesIO
import torch
from transformers import AutoModelForImageClassification, AutoImageProcessor


app = FastAPI()

model_name = "prithivMLmods/open-deepfake-detection"
extractor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)

@app.post("/analyze")
async def analyze(request: Request):
    data = await request.json()
    image_url = data.get("mediaUrl")
    if not image_url:
        return {"error": "Missing 'mediaUrl'"}

    try:
        img_bytes = requests.get(image_url).content
        img = Image.open(BytesIO(img_bytes)).convert("RGB")
        inputs = extractor(images=img, return_tensors="pt")
        with torch.no_grad():
            outputs = model(**inputs)
        scores = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
        confidence, pred_idx = torch.max(scores, dim=0)
        label = model.config.id2label[pred_idx.item()]

        return {
            "label": label.lower(),
            "score": round(confidence.item(), 3)
        }

    except Exception as e:
        return {"error": str(e)}