Putting It All Together – AI Video Analysis System

from fastapi import FastAPI, WebSocket, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
import asyncio, httpx, subprocess, uuid, os, json, tempfile

app = FastAPI(title="AI Video Analysis System — Full Demo")

jobs: dict = {}  # job_id → {status, results}

# ── REST: 비디오 업로드 & 분석 시작 ─────────────────────────
@app.post("/analyze")
async def start_analysis(video: UploadFile = File(...),
                         prompt: str = "이 영상을 설명해주세요"):
    job_id = str(uuid.uuid4())
    jobs[job_id] = {"status": "queued", "results": [], "prompt": prompt}

    # 비디오 파일을 임시 저장
    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp:
        tmp.write(await video.read())
        tmp_path = tmp.name

    # 백그라운드에서 분석 시작 (즉시 job_id 반환)
    asyncio.create_task(_analyze(job_id, tmp_path, prompt))
    return {"job_id": job_id, "status": "queued",
            "ws_url": f"ws://localhost:8000/ws/{job_id}"}

# ── WebSocket: 실시간 진행 상황 스트리밍 ────────────────────
@app.websocket("/ws/{job_id}")
async def progress_stream(ws: WebSocket, job_id: str):
    await ws.accept()
    try:
        while True:
            job = jobs.get(job_id)
            if not job:
                await ws.send_json({"error": "job not found"}); break
            await ws.send_json(job)          # 현재 상태 전송
            if job["status"] in ("done", "error"):
                break
            await asyncio.sleep(0.5)         # 0.5초마다 업데이트
    except Exception:
        pass

# ── 내부: ffmpeg → Worker 병렬 호출 ─────────────────────────
async def _analyze(job_id: str, video_path: str, prompt: str):
    try:
        # 1) ffmpeg로 1초마다 프레임 추출
        jobs[job_id]["status"] = "extracting_frames"
        frame_dir = tempfile.mkdtemp()
        subprocess.run(
            ["ffmpeg", "-i", video_path, "-vf", "fps=1",
             f"{frame_dir}/frame_%03d.jpg", "-y", "-loglevel", "quiet"],
            check=True
        )
        frame_count = len(os.listdir(frame_dir))
        jobs[job_id]["frame_count"] = frame_count

        # 2) 3개 워커에 동시 요청 (챕터 5 asyncio.gather 패턴)
        jobs[job_id]["status"] = "analyzing"
        worker_urls = [
            "http://worker-a:8001/infer",   # Qwen 27B
            "http://worker-b:8002/infer",   # Qwen 9B
            "http://worker-c:8003/infer",   # Qwen 4B
        ]
        async with httpx.AsyncClient(timeout=120) as client:
            tasks = [
                client.post(url, json={"prompt": prompt,
                                       "frame_dir": frame_dir})
                for url in worker_urls
            ]
            responses = await asyncio.gather(*tasks, return_exceptions=True)

        results = []
        for i, r in enumerate(responses):
            if isinstance(r, Exception):
                results.append({"worker": i+1, "error": str(r)})
            else:
                results.append({"worker": i+1, "answer": r.json()["answer"]})

        jobs[job_id].update({"status": "done", "results": results})

    except Exception as e:
        jobs[job_id].update({"status": "error", "error": str(e)})
    finally:
        os.unlink(video_path)   # 임시 파일 정리

# 실행: uvicorn main:app --reload --port 8000