app.py

from fastapi import FastAPI, UploadFile, File, Response,Header, BackgroundTasks,Body
from fastapi.staticfiles import StaticFiles
from vitpose import VitPose
import os 
from dotenv import load_dotenv
from tasks import process_video
from fastapi.responses import JSONResponse
from config import API_KEY
import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
app = FastAPI()
vitpose = VitPose()
# vitpose.pipeline.warmup()

load_dotenv()


app.mount("/static", StaticFiles())

@app.get("/")
def read_root():
    return {"message": "Hello World"}

@app.get("/test")
def test():
    return {"message": "from test"}

@app.post("/upload")
async def upload(background_tasks: BackgroundTasks,
                 file: UploadFile = File(...), 
                 token: str = Header(...),
                 user_id: str = Body(...)):
    
    logger.info(f"Uploading video {file.filename} {token} {user_id}")
    
    if token != API_KEY:
        logger.info(f"Unauthorized {token} {API_KEY}")
        return JSONResponse(content={"message": "Unauthorized", "status": 401})
    
    logger.info("reading contents")
    contents = await file.read()
    
    # Save the file to the local directory
    logger.info("saving file")
    with open(file.filename, "wb") as f:
        f.write(contents)
    
    logger.info(f"file saved {file.filename}")
    
    # Create a clone of the file with content already read
    background_tasks.add_task(process_video, file.filename, vitpose, user_id)  

    # Return the file as a response
    return JSONResponse(content={"message": "Video uploaded successfully", "status": 200})