deepcheck-api / app.py
imiksi's picture
Update app.py
b12f359 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, AutoFeatureExtractor
app = FastAPI()
model_name = "prithivMLmods/open-deepfake-detection"
extractor = AutoFeatureExtractor.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)}