Upload app.py

app.py

@@ -1,23 +1,45 @@
 import os
-import json
+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
 
-def run_pipeline():
+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"
-    frames_dir = "frames"
     cursor_dir = "cursors"
-    annotations_dir = "annotations"
-    os.makedirs(frames_dir, exist_ok=True)
+    os.makedirs(FRAMES_DIR, exist_ok=True)
     os.makedirs(cursor_dir, exist_ok=True)
-    os.makedirs(annotations_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)
+            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"))
-    print("Pipeline complete.")
+    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)
 
-if __name__ == "__main__":
-    run_pipeline()
+@app.get("/download/annotations")
+def download_annotations():
+    return FileResponse(ANNOTATIONS_ZIP, filename=ANNOTATIONS_ZIP)