dappai's picture
Upload 9 files
459fc8b verified
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
import shutil
import uuid
import os
from inference import run_inference, generate_heatmap
app = FastAPI()
@app.get("/")
def root():
return {"status": "Deepfake API running"}
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
@app.post("/predict")
async def predict(video: UploadFile = File(...)):
file_id = f"{uuid.uuid4()}.mp4"
path = f"/tmp/{file_id}"
try:
with open(path, "wb") as buffer:
shutil.copyfileobj(video.file, buffer)
result = run_inference(path)
return result
except Exception as e:
raise HTTPException(
status_code=500,
detail=str(e)
)
finally:
if os.path.exists(path):
os.remove(path)
@app.post("/heatmap")
async def heatmap(data: dict):
try:
frame_index = int(data.get("frame_index", -1))
if frame_index < 0:
raise HTTPException(
status_code=400,
detail="Invalid frame index"
)
heatmap, regions = generate_heatmap(frame_index)
if heatmap is None:
raise HTTPException(
status_code=404,
detail="Frame not found"
)
return {
"heatmap": heatmap,
"regions": regions
}
except Exception as e:
raise HTTPException(
status_code=500,
detail=str(e)
)