| 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) |
| ) |