Update app.py

app.py

@@ -1,45 +1,48 @@
-import os
-import zipfile
-from fastapi import FastAPI
-from fastapi.responses import FileResponse
-from scripts.extract_frames import extract_frames
-from scripts.cursor_tracker import track_cursor
-
-app = FastAPI()
-
-FRAMES_DIR = "frames"
-ANNOTATIONS_DIR = "annotations"
-FRAMES_ZIP = "frames.zip"
-ANNOTATIONS_ZIP = "annotations.zip"
-
-@app.on_event("startup")
-def run_pipeline_and_zip():
-    video_dir = "videos"
-    cursor_dir = "cursors"
-    os.makedirs(FRAMES_DIR, exist_ok=True)
-    os.makedirs(cursor_dir, exist_ok=True)
-    os.makedirs(ANNOTATIONS_DIR, exist_ok=True)
-    # Step 1: Extract frames
-    for video_file in os.listdir(video_dir):
-        if video_file.lower().endswith(('.mp4', '.avi', '.mov')):
-            extract_frames(os.path.join(video_dir, video_file), FRAMES_DIR)
-    # Step 2: Track cursor
-    track_cursor(FRAMES_DIR, cursor_dir, os.path.join(ANNOTATIONS_DIR, "cursor_positions.json"))
-    # Step 3: Zip frames and annotations
-    with zipfile.ZipFile(FRAMES_ZIP, 'w', zipfile.ZIP_DEFLATED) as zipf:
-        for root, _, files in os.walk(FRAMES_DIR):
-            for file in files:
-                zipf.write(os.path.join(root, file), arcname=file)
-    with zipfile.ZipFile(ANNOTATIONS_ZIP, 'w', zipfile.ZIP_DEFLATED) as zipf:
-        for root, _, files in os.walk(ANNOTATIONS_DIR):
-            for file in files:
-                zipf.write(os.path.join(root, file), arcname=file)
-    print("Pipeline complete and outputs zipped.")
-
-@app.get("/download/frames")
-def download_frames():
-    return FileResponse(FRAMES_ZIP, filename=FRAMES_ZIP)
-
-@app.get("/download/annotations")
-def download_annotations():
-    return FileResponse(ANNOTATIONS_ZIP, filename=ANNOTATIONS_ZIP)
+import os
+import zipfile
+from fastapi import FastAPI
+from fastapi.responses import FileResponse
+from scripts.extract_frames import extract_frames
+from scripts.cursor_tracker import track_cursor
+from contextlib import asynccontextmanager
+
+FRAMES_DIR = "frames"
+ANNOTATIONS_DIR = "annotations"
+FRAMES_ZIP = "frames.zip"
+ANNOTATIONS_ZIP = "annotations.zip"
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # --- This replaces the old @app.on_event("startup") ---
+    video_dir = "videos"
+    cursor_dir = "cursors"
+    os.makedirs(FRAMES_DIR, exist_ok=True)
+    os.makedirs(cursor_dir, exist_ok=True)
+    os.makedirs(ANNOTATIONS_DIR, exist_ok=True)
+    # Step 1: Extract frames
+    for video_file in os.listdir(video_dir):
+        if video_file.lower().endswith(('.mp4', '.avi', '.mov')):
+            extract_frames(os.path.join(video_dir, video_file), FRAMES_DIR)
+    # Step 2: Track cursor
+    track_cursor(FRAMES_DIR, cursor_dir, os.path.join(ANNOTATIONS_DIR, "cursor_positions.json"))
+    # Step 3: Zip frames and annotations
+    with zipfile.ZipFile(FRAMES_ZIP, 'w', zipfile.ZIP_DEFLATED) as zipf:
+        for root, _, files in os.walk(FRAMES_DIR):
+            for file in files:
+                zipf.write(os.path.join(root, file), arcname=file)
+    with zipfile.ZipFile(ANNOTATIONS_ZIP, 'w', zipfile.ZIP_DEFLATED) as zipf:
+        for root, _, files in os.walk(ANNOTATIONS_DIR):
+            for file in files:
+                zipf.write(os.path.join(root, file), arcname=file)
+    print("Pipeline complete and outputs zipped.")
+    yield  # This is required for the lifespan context
+
+app = FastAPI(lifespan=lifespan)
+
+@app.get("/download/frames")
+def download_frames():
+    return FileResponse(FRAMES_ZIP, filename=FRAMES_ZIP)
+
+@app.get("/download/annotations")
+def download_annotations():
+    return FileResponse(ANNOTATIONS_ZIP, filename=ANNOTATIONS_ZIP)