Redfire-1234's picture
Update backend/app.py
f06644b verified
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, FileResponse
from fastapi.staticfiles import StaticFiles
import os
import shutil
from pathlib import Path
import json
from tracker import VideoTracker
from utils import save_trajectory_json, create_output_dirs
app = FastAPI(title="3D Trajectory Tracker")
# CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Create necessary directories
BASE_DIR = Path(__file__).parent.parent
UPLOAD_DIR = BASE_DIR / "uploads"
OUTPUT_DIR = BASE_DIR / "outputs"
FRONTEND_DIR = BASE_DIR / "frontend"
create_output_dirs([UPLOAD_DIR, OUTPUT_DIR])
# Mount static files
if FRONTEND_DIR.exists():
app.mount("/static", StaticFiles(directory=FRONTEND_DIR), name="static")
@app.get("/")
async def read_root():
"""Serve the main HTML page"""
index_path = FRONTEND_DIR / "index.html"
if index_path.exists():
return FileResponse(index_path)
return {"message": "3D Trajectory Tracker API"}
@app.post("/api/upload")
async def upload_video(file: UploadFile = File(...)):
"""Upload and process video file"""
try:
# Validate file type
if not file.filename.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
raise HTTPException(400, "Invalid file type. Please upload a video file.")
# Save uploaded file
file_path = UPLOAD_DIR / file.filename
with open(file_path, "wb") as buffer:
shutil.copyfileobj(file.file, buffer)
# Process video
tracker = VideoTracker(str(file_path))
trajectories = tracker.process_video()
# Save results
output_path = OUTPUT_DIR / f"{Path(file.filename).stem}_trajectories.json"
save_trajectory_json(trajectories, output_path)
# Clean up uploaded file
os.remove(file_path)
return JSONResponse({
"success": True,
"trajectories": trajectories,
"output_file": str(output_path)
})
except Exception as e:
raise HTTPException(500, f"Processing error: {str(e)}")
@app.get("/api/trajectories/{filename}")
async def get_trajectory(filename: str):
"""Retrieve saved trajectory data"""
file_path = OUTPUT_DIR / filename
if not file_path.exists():
raise HTTPException(404, "Trajectory file not found")
with open(file_path, "r") as f:
data = json.load(f)
return JSONResponse(data)
@app.get("/api/list")
async def list_trajectories():
"""List all processed trajectories"""
files = [f.name for f in OUTPUT_DIR.glob("*.json")]
return JSONResponse({"files": files})
@app.delete("/api/clear")
async def clear_data():
"""Clear all uploaded and processed files"""
for folder in [UPLOAD_DIR, OUTPUT_DIR]:
for file in folder.glob("*"):
if file.is_file():
os.remove(file)
return JSONResponse({"success": True, "message": "All data cleared"})
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)