Update streamlit_app.py

streamlit_app.py

@@ -11,6 +11,7 @@
 import traceback
 import uuid
 from tempfile import NamedTemporaryFile
+import threading
 import streamlit as st
 
 from ui import render_ui
@@ -21,6 +22,7 @@
     check_gpu
 )
 from models.model_loaders import load_sam2, load_matanyone
+from utils.progress_tracker import init_progress, update_progress, mark_complete, get_progress
 
 APP_NAME = "Advanced Video Background Replacer"
 LOG_FILE = "/tmp/app.log"
@@ -31,7 +33,6 @@ def setup_logging(level: int = logging.INFO) -> logging.Logger:
     logger = logging.getLogger(APP_NAME)
     logger.setLevel(level)
     logger.propagate = False
-    # Remove previous handlers (Streamlit reruns)
     for h in list(logger.handlers):
         logger.removeHandler(h)
     ch = logging.StreamHandler(sys.stdout)
@@ -52,7 +53,6 @@ def custom_excepthook(type, value, tb):
     logger.error(f"Unhandled: {type.__name__}: {value}\n{''.join(traceback.format_tb(tb))}", exc_info=True)
 sys.excepthook = custom_excepthook
 
-# Only load once
 sam2_predictor = load_sam2()
 matanyone_processor = load_matanyone()
 
@@ -83,19 +83,17 @@ 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, progress_callback=None):
+def process_video_background(uploaded_video, background, bg_type):
+    """Background thread for video processing"""
     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())}")
-    logger.info(f"[RUN {run_id}] Video size={len(uploaded_video.read()) / 1e6:.2f}MB, BG type={bg_type}")
-    uploaded_video.seek(0)
-    st.session_state.processing = True
-    st.session_state.processed_video_bytes = None
-    st.session_state.last_error = None
+    
     t0 = time.time()
     try:
-        if progress_callback:
-            progress_callback("📥 Uploading video...")
+        init_progress()
+        update_progress("Starting video processing...", 0, "Initialization")
+        
         suffix = Path(uploaded_video.name).suffix or ".mp4"
         with NamedTemporaryFile(delete=False, suffix=suffix) as tmp_vid:
             uploaded_video.seek(0)
@@ -103,49 +101,81 @@ def process_video(uploaded_video, background, bg_type, progress_callback=None):
             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)...")
-        # >>> CHANGE: pass a short watchdog timeout so hangs surface quickly
+        def progress_cb(msg):
+            # Map messages to progress percentages and stages
+            progress = 0
+            stage = "Processing"
+            
+            if "Stage 1 initiated" in msg:
+                progress, stage = 5, "Stage 1"
+            elif "GPU engaged" in msg:
+                progress, stage = 10, "Stage 1 - SAM2"
+            elif "SAM2 processing frame" in msg:
+                progress, stage = 15, "Stage 1 - SAM2"
+            elif "SAM2 complete" in msg:
+                progress, stage = 25, "Stage 1 - SAM2"
+            elif "MatAnyone starting" in msg:
+                progress, stage = 30, "Stage 1 - MatAnyone"
+            elif "MatAnyone processing" in msg:
+                progress, stage = 50, "Stage 1 - MatAnyone"
+            elif "MatAnyone complete" in msg:
+                progress, stage = 70, "Stage 1 - MatAnyone"
+            elif "Smoothing" in msg:
+                progress, stage = 75, "Stage 1 - Smoothing"
+            elif "transparent video" in msg:
+                progress, stage = 80, "Stage 1 - Finalizing"
+            elif "Stage 1 complete" in msg:
+                progress, stage = 85, "Stage 1 Complete"
+            elif "Stage 2 begun" in msg:
+                progress, stage = 86, "Stage 2"
+            elif "Compositing" in msg:
+                progress, stage = 90, "Stage 2 - Compositing"
+            elif "Restoring audio" in msg:
+                progress, stage = 95, "Stage 2 - Audio"
+            elif "Stage 2 complete" in msg:
+                progress, stage = 98, "Stage 2 Complete"
+                
+            update_progress(msg, progress, stage)
+            logger.info(f"[PROGRESS] {msg}")
+
         transparent_path, audio_path = stage1_create_transparent_video(
             tmp_vid_path,
             sam2_predictor=sam2_predictor,
             matanyone_processor=matanyone_processor,
-            mat_timeout_sec=300
+            mat_timeout_sec=300,
+            progress_callback=progress_cb
         )
+        
         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}")
+        logger.info(f"[RUN {run_id}] Stage 1 completed")
 
-        if progress_callback:
-            progress_callback("🎨 Stage 2: Compositing with background and restoring audio...")
         final_path = stage2_composite_background(
             transparent_path,
             audio_path,
             background,
-            bg_type.lower()
+            bg_type.lower(),
+            progress_callback=progress_cb
         )
+        
         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}")
+        logger.info(f"[RUN {run_id}] Stage 2 completed")
 
-        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
+        logger.info(f"[RUN {run_id}] SUCCESS total Δ={total:.2f}s")
+        mark_complete(success=True)
+        
     except Exception as e:
         total = time.time() - t0
-        error_msg = f"[RUN {run_id}] Processing Error: {str(e)} (Δ {total:.2f}s)\n\nCheck logs for details."
+        error_msg = f"Processing Error: {str(e)} (Δ {total:.2f}s)"
         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
+        mark_complete(success=False, error=str(e))
     finally:
         st.session_state.processing = False
         logger.info(f"[RUN {run_id}] Processing finished")
@@ -159,12 +189,11 @@ def main():
             initial_sidebar_state="expanded"
         )
         initialize_session_state()
-        render_ui(process_video)
+        render_ui(process_video_background)
     except Exception as e:
         logger.error(f"Main app error: {e}", exc_info=True)
         st.error(f"App startup failed: {str(e)}. Check logs for details.")
 
 if __name__ == "__main__":
     setup_t4_environment()
-    main()
- 
+    main()