Update streamlit_app.py

streamlit_app.py

@@ -83,7 +83,7 @@ def initialize_session_state():
     if st.session_state.gpu_available is None:
         st.session_state.gpu_available = check_gpu(logger)
 
-def process_video(uploaded_video, background, bg_type):
+def process_video(uploaded_video, background, bg_type, progress_callback=None):
     run_id = uuid.uuid4().hex[:8]
     logger.info("=" * 80)
     logger.info(f"[RUN {run_id}] VIDEO PROCESSING STARTED at {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())}")
@@ -94,12 +94,17 @@ def process_video(uploaded_video, background, bg_type):
     st.session_state.last_error = None
     t0 = time.time()
     try:
+        if progress_callback:
+            progress_callback("📥 Uploading video...")
         suffix = Path(uploaded_video.name).suffix or ".mp4"
         with NamedTemporaryFile(delete=False, suffix=suffix) as tmp_vid:
             uploaded_video.seek(0)
             tmp_vid.write(uploaded_video.read())
             tmp_vid_path = tmp_vid.name
         logger.info(f"[RUN {run_id}] Temporary video path: {tmp_vid_path}")
+
+        if progress_callback:
+            progress_callback("🚀 Stage 1: Creating transparent video (matting & segmentation)...")
         transparent_path, audio_path = stage1_create_transparent_video(
             tmp_vid_path,
             sam2_predictor=sam2_predictor,
@@ -108,7 +113,9 @@ def process_video(uploaded_video, background, bg_type):
         if not transparent_path or not os.path.exists(transparent_path):
             raise RuntimeError("Stage 1 failed: Transparent video not created")
         logger.info(f"[RUN {run_id}] Stage 1 completed: Transparent path={transparent_path}, Audio path={audio_path}")
-        # For "color", background is "#RRGGBB", for others it's PIL.Image
+
+        if progress_callback:
+            progress_callback("🎨 Stage 2: Compositing with background and restoring audio...")
         final_path = stage2_composite_background(
             transparent_path,
             audio_path,
@@ -118,10 +125,15 @@ def process_video(uploaded_video, background, bg_type):
         if not final_path or not os.path.exists(final_path):
             raise RuntimeError("Stage 2 failed: Final video not created")
         logger.info(f"[RUN {run_id}] Stage 2 completed: Final path={final_path}")
+
+        if progress_callback:
+            progress_callback("📤 Loading final video for download...")
         with open(final_path, 'rb') as f:
             st.session_state.processed_video_bytes = f.read()
         total = time.time() - t0
         logger.info(f"[RUN {run_id}] SUCCESS size={len(st.session_state.processed_video_bytes)/1e6:.2f}MB, total Δ={total:.2f}s")
+        if progress_callback:
+            progress_callback("✅ All done!")
         return True
     except Exception as e:
         total = time.time() - t0
@@ -129,6 +141,8 @@ def process_video(uploaded_video, background, bg_type):
         logger.error(error_msg)
         logger.error(traceback.format_exc())
         st.session_state.last_error = error_msg
+        if progress_callback:
+            progress_callback(f"❌ ERROR: {str(e)}")
         return False
     finally:
         st.session_state.processing = False