dappai's picture
Upload folder using huggingface_hub
4a9bdfe verified
Raw
History Blame Contribute Delete
1.64 kB
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)
)