Update pipeline_full.py

pipeline_full.py

@@ -1,5 +1,6 @@
 # pipeline_full.py
 import os
+import json
 import base64
 from io import BytesIO
 from typing import List, Dict, Any
@@ -30,7 +31,7 @@ from sports.annotators.soccer import (
 )
 
 # ------------------------------------------------------------------
-# Globals – will be initialized lazily so build/startup doesn't crash
+# Globals – initialized lazily so build/startup doesn't crash
 # ------------------------------------------------------------------
 
 PLAYER_DETECTION_MODEL = None
@@ -47,6 +48,31 @@ REFEREE_ID = 3
 
 MODELS_READY = False
 
+# progress tracking
+CURRENT_JOB_DIR: str | None = None
+
+
+def set_job_dir(job_dir: str):
+    global CURRENT_JOB_DIR
+    CURRENT_JOB_DIR = job_dir
+
+
+def update_progress(stage: str, progress: float, message: str = ""):
+    """
+    Write a small JSON status file in the current job dir so the UI can poll.
+    """
+    if not CURRENT_JOB_DIR:
+        return
+    status = {
+        "stage": stage,
+        "progress": float(progress),
+        "message": message,
+    }
+    os.makedirs(CURRENT_JOB_DIR, exist_ok=True)
+    status_path = os.path.join(CURRENT_JOB_DIR, "status.json")
+    with open(status_path, "w", encoding="utf-8") as f:
+        json.dump(status, f)
+
 
 def ensure_models_loaded():
     """
@@ -80,7 +106,7 @@ def ensure_models_loaded():
 
     # SigLIP embeddings
     SIGLIP_MODEL_PATH = "google/siglip-base-patch16-224"
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = get_device()
     EMBEDDINGS_MODEL = SiglipVisionModel.from_pretrained(SIGLIP_MODEL_PATH).to(device)
     EMBEDDINGS_PROCESSOR = AutoProcessor.from_pretrained(SIGLIP_MODEL_PATH)
 
@@ -370,7 +396,7 @@ def resolve_goalkeepers_team_id(players: sv.Detections, goalkeepers: sv.Detectio
     return np.array(goalkeepers_team_id)
 
 
-# -------------------- 5. Voronoi blend helper (your function) --------------------
+# -------------------- 5. Voronoi blend helper --------------------
 
 
 def draw_pitch_voronoi_diagram_2(
@@ -747,21 +773,43 @@ def run_full_pipeline(video_path: str, job_dir: str) -> Dict[str, Any]:
     Run the full notebook-equivalent pipeline on a video and save all artifacts
     into job_dir. Returns paths + stats for the FastAPI app.
     """
+    set_job_dir(job_dir)
+    update_progress("initializing", 0.0, "Loading models...")
     ensure_models_loaded()
+
     os.makedirs(job_dir, exist_ok=True)
 
+    update_progress("siglip", 0.15, "Running SigLIP clustering...")
     siglip_out = step_siglip_clustering(video_path, os.path.join(job_dir, "siglip"))
+
+    update_progress("team_classifier", 0.30, "Training TeamClassifier...")
     train_team_classifier_on_video(video_path)
 
+    update_progress("basic_frames", 0.45, "Generating basic annotated frames...")
     basic_paths = step_basic_frames(video_path, os.path.join(job_dir, "frames"))
+
+    update_progress("advanced_views", 0.60, "Generating advanced radar / Voronoi views...")
     adv_paths = step_single_frame_advanced(video_path, os.path.join(job_dir, "advanced"))
+
+    update_progress("ball_path", 0.80, "Computing ball path and heatmap...")
     ball_paths = step_ball_path(video_path, os.path.join(job_dir, "ball_path"))
+
+    update_progress("stats", 0.90, "Calculating stats...")
     stats = process_video_stats(video_path)
 
-    return {
+    result = {
         "basic": basic_paths,
         "advanced": adv_paths,
         "ball": ball_paths,
         "stats": stats,
         "siglip_html": siglip_out["plot_html"],
     }
+
+    # Save a copy for the UI result page
+    result_path = os.path.join(job_dir, "result.json")
+    with open(result_path, "w", encoding="utf-8") as f:
+        json.dump(result, f)
+
+    update_progress("done", 1.0, "Completed")
+
+    return result