deepcheck-api / app.py
imiksi's picture
Update app.py
95ab217 verified
raw
history blame
1.19 kB
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)}