Update streamlit_app.py

streamlit_app.py

@@ -14,6 +14,7 @@
 import io
 import torch
 import traceback
+import shutil
 
 # --- Project Setup ---
 sys.path.append(str(Path(__file__).parent.absolute()))
@@ -88,20 +89,19 @@ def initialize_session_state():
     defaults = {
         'uploaded_video': None,
         'video_bytes_cache': None,
-        'video_preview_placeholder': None,
         'bg_image': None,
         'bg_image_cache': None,
         'bg_image_name': None,
-        'bg_preview_placeholder': None,
         'bg_color': "#00FF00",
         'cached_color': None,
         'color_display_cache': None,
-        'processed_video_path': None,
+        'processed_video_bytes': None,  # Store bytes instead of path
         'processing': False,
         'progress': 0,
         'progress_text': "Ready",
         'last_video_id': None,
-        'last_bg_image_id': None
+        'last_bg_image_id': None,
+        'process_complete': False
     }
     for key, value in defaults.items():
         if key not in st.session_state:
@@ -111,93 +111,148 @@ def initialize_session_state():
 def process_video(input_file, background, bg_type="image"):
     """
     Process video with the selected background using SAM2 and MatAnyone pipeline.
-    Returns the path to the processed video file.
+    Returns the video bytes.
     """
+    logger.info("=" * 60)
+    logger.info("🎬 STARTING VIDEO PROCESSING")
+    logger.info("=" * 60)
+    
     try:
-        logger.info("==== Starting process_video ====")
-        # Create a temporary directory for processing
-        with tempfile.TemporaryDirectory() as temp_dir:
-            temp_dir = Path(temp_dir)
-            logger.info(f"Initialized temp_dir: {temp_dir}")
+        # Create a persistent temporary directory
+        temp_base = Path(tempfile.gettempdir()) / "video_processing"
+        temp_base.mkdir(exist_ok=True)
+        temp_dir = temp_base / f"session_{int(time.time())}"
+        temp_dir.mkdir(exist_ok=True)
+        
+        logger.info(f"📁 Created temp directory: {temp_dir}")
+
+        # Save the uploaded video to a temporary file
+        input_path = str(temp_dir / "input.mp4")
+        logger.info(f"💾 Writing video to {input_path}")
+        
+        with open(input_path, "wb") as f:
+            written = f.write(input_file.getvalue())
+            logger.info(f"✅ Wrote {written/1e6:.2f}MB")
+        
+        if not os.path.exists(input_path):
+            raise FileNotFoundError(f"❌ Input video not saved: {input_path}")
+        
+        logger.info(f"✅ Video file exists: {os.path.getsize(input_path)} bytes")
+
+        # Prepare background
+        bg_path = None
+        if bg_type == "image" and background is not None:
+            logger.info("🖼️ Processing background IMAGE")
+            bg_cv = cv2.cvtColor(np.array(background), cv2.COLOR_RGB2BGR)
+            bg_path = str(temp_dir / "background.jpg")
+            cv2.imwrite(bg_path, bg_cv)
+            logger.info(f"✅ Background image written to {bg_path}")
+            
+        elif bg_type == "color" and hasattr(st.session_state, 'bg_color'):
+            logger.info(f"🎨 Processing background COLOR: {st.session_state.bg_color}")
+            color_hex = st.session_state.bg_color.lstrip('#')
+            color_rgb = tuple(int(color_hex[i:i+2], 16) for i in (0, 2, 4))
+            bg_path = str(temp_dir / "background.jpg")
+            cv2.imwrite(bg_path, np.ones((100, 100, 3), dtype=np.uint8) * color_rgb[::-1])
+            logger.info(f"✅ Background color image written to {bg_path}")
+
+        # Set up progress placeholders
+        progress_placeholder = st.empty()
+        status_placeholder = st.empty()
+        
+        def progress_callback(progress, message):
+            progress = max(0, min(1, float(progress)))
+            logger.info(f"📊 Progress: {progress*100:.1f}% - {message}")
+            progress_placeholder.progress(progress)
+            status_placeholder.text(f"Status: {message}")
 
-            # Save the uploaded video to a temporary file
-            input_path = str(temp_dir / "input.mp4")
-            logger.info(f"Writing video to {input_path}")
-            with open(input_path, "wb") as f:
-                written = f.write(input_file.getvalue())
-                logger.info(f"Wrote {written/1e6:.2f}MB")
-            if not os.path.exists(input_path):
-                raise FileNotFoundError(f"Input video not saved: {input_path}")
+        # Log GPU state
+        if torch.cuda.is_available():
+            logger.info(f"🎮 CUDA Available: True, Device: {torch.cuda.get_device_name(0)}")
+            logger.info(f"💾 GPU Mem Before: {torch.cuda.memory_allocated()/1e9:.2f}GB")
+        else:
+            logger.info("⚠️ CUDA Available: False (using CPU)")
 
-            # Prepare background
-            bg_path = None
-            if bg_type == "image" and background is not None:
-                bg_cv = cv2.cvtColor(np.array(background), cv2.COLOR_RGB2BGR)
-                bg_path = str(temp_dir / "background.jpg")
-                cv2.imwrite(bg_path, bg_cv)
-                logger.info(f"Background image written to {bg_path}")
-            elif bg_type == "color" and hasattr(st.session_state, 'bg_color'):
-                color_hex = st.session_state.bg_color.lstrip('#')
-                color_rgb = tuple(int(color_hex[i:i+2], 16) for i in (0, 2, 4))
-                bg_path = str(temp_dir / "background.jpg")
-                cv2.imwrite(bg_path, np.ones((100, 100, 3), dtype=np.uint8) * color_rgb[::-1])
-                logger.info(f"Background color image written to {bg_path}")
+        # Process the video
+        output_path = str(temp_dir / "output.mp4")
+        click_points = [[0.5, 0.5]]
+        
+        logger.info("🔧 Importing TwoStageProcessor...")
+        from pipeline.integrated_pipeline import TwoStageProcessor
 
-            logger.info(f"Disk free before processing: {os.popen('df -h / | tail -1').read().strip()}")
+        # Cache processor
+        @st.cache_resource
+        def load_processor(temp_dir_str):
+            logger.info(f"🏗️ Loading processor with temp_dir: {temp_dir_str}")
+            return TwoStageProcessor(temp_dir=temp_dir_str)
+        
+        processor = load_processor(str(temp_dir))
+        logger.info("✅ Processor loaded")
+
+        try:
+            logger.info("🎬 Calling TwoStageProcessor.process_video...")
+            logger.info(f"   - Input: {input_path}")
+            logger.info(f"   - Background: {bg_path}")
+            logger.info(f"   - Output: {output_path}")
+            logger.info(f"   - Click points: {click_points}")
+            
+            success = processor.process_video(
+                input_video=input_path,
+                background_video=bg_path if bg_type == "image" else "",
+                click_points=click_points,
+                output_path=output_path,
+                use_matanyone=True,
+                progress_callback=progress_callback
+            )
             
-            # Set up progress placeholders
-            progress_placeholder = st.empty()
-            status_placeholder = st.empty()
+            logger.info(f"📊 Processing returned: {success}")
             
-            def progress_callback(progress, message):
-                progress = max(0, min(1, float(progress)))
-                progress_placeholder.progress(progress)
-                status_placeholder.text(f"Status: {message}")
-
-            # Log GPU state
-            if torch.cuda.is_available():
-                logger.info(f"CUDA Available: True, Device: {torch.cuda.get_device_name(0)}")
-                logger.info(f"GPU Mem Before: {torch.cuda.memory_allocated()/1e9:.2f}GB")
-            else:
-                logger.info("CUDA Available: False")
+        except Exception as e:
+            logger.error(f"❌ Pipeline processing failed: {traceback.format_exc()}", exc_info=True)
+            raise
 
-            # Process the video
-            output_path = str(temp_dir / "output.mp4")
-            click_points = [[0.5, 0.5]]
-            from pipeline.integrated_pipeline import TwoStageProcessor
+        if torch.cuda.is_available():
+            logger.info(f"💾 GPU Mem After: {torch.cuda.memory_allocated()/1e9:.2f}GB")
+            torch.cuda.empty_cache()
+            logger.info("🧹 GPU cache cleared")
 
-            # Cache processor
-            @st.cache_resource
-            def load_processor(temp_dir):
-                return TwoStageProcessor(temp_dir=temp_dir)
-            processor = load_processor(str(temp_dir))
+        if not success:
+            raise RuntimeError("❌ Video processing returned False")
 
-            try:
-                logger.info("Calling TwoStageProcessor.process_video...")
-                success = processor.process_video(
-                    input_video=input_path,
-                    background_video=bg_path if bg_type == "image" else "",
-                    click_points=click_points,
-                    output_path=output_path,
-                    use_matanyone=True,
-                    progress_callback=progress_callback
-                )
-            except Exception as e:
-                logger.error(f"Pipeline processing failed: {traceback.format_exc()}", exc_info=True)
-                raise
+        # Check if output file exists
+        if not os.path.exists(output_path):
+            logger.error(f"❌ Output file does not exist: {output_path}")
+            raise FileNotFoundError(f"Output video not created: {output_path}")
+        
+        output_size = os.path.getsize(output_path)
+        logger.info(f"✅ Output file exists: {output_size} bytes ({output_size/1e6:.2f}MB)")
 
-            if torch.cuda.is_available():
-                logger.info(f"GPU Mem After: {torch.cuda.memory_allocated()/1e9:.2f}GB")
-                torch.cuda.empty_cache()
+        # Read the output video into memory
+        logger.info("📖 Reading output video into memory...")
+        with open(output_path, 'rb') as f:
+            video_bytes = f.read()
+        
+        logger.info(f"✅ Read {len(video_bytes)/1e6:.2f}MB into memory")
 
-            if not success:
-                raise RuntimeError("Video processing failed")
+        # Clean up temp directory
+        try:
+            logger.info(f"🧹 Cleaning up temp directory: {temp_dir}")
+            shutil.rmtree(temp_dir)
+            logger.info("✅ Temp directory cleaned")
+        except Exception as e:
+            logger.warning(f"⚠️ Could not clean temp directory: {e}")
 
-            logger.info("==== process_video completed successfully ====")
-            return output_path
+        logger.info("=" * 60)
+        logger.info("✅ VIDEO PROCESSING COMPLETED SUCCESSFULLY")
+        logger.info("=" * 60)
+        
+        return video_bytes
 
     except Exception as e:
-        logger.error(f"Error in video processing: {str(e)}", exc_info=True)
+        logger.error("=" * 60)
+        logger.error(f"❌ ERROR IN VIDEO PROCESSING: {str(e)}")
+        logger.error(traceback.format_exc())
+        logger.error("=" * 60)
         st.error(f"An error occurred during processing: {str(e)}")
         return None
 
@@ -209,6 +264,8 @@ def main():
     # Initialize session state
     initialize_session_state()
     
+    logger.info(f"🔄 App rerun - Processing: {st.session_state.processing}, Complete: {st.session_state.process_complete}")
+    
     # Main layout
     col1, col2 = st.columns([1, 1], gap="large")
     
@@ -224,24 +281,28 @@ def main():
         # Check if video actually changed using id()
         current_video_id = id(uploaded)
         if current_video_id != st.session_state.last_video_id:
+            logger.info(f"📹 New video uploaded: {uploaded.name if uploaded else 'None'}")
             st.session_state.uploaded_video = uploaded
             st.session_state.last_video_id = current_video_id
             st.session_state.video_bytes_cache = None
+            st.session_state.processed_video_bytes = None  # Clear processed video
+            st.session_state.process_complete = False
         
         # Video preview section
         st.markdown("### Video Preview")
-        if st.session_state.video_preview_placeholder is None:
-            st.session_state.video_preview_placeholder = st.empty()
+        video_preview_container = st.container()
         
-        if st.session_state.uploaded_video is not None:
-            if st.session_state.video_bytes_cache is None:
-                st.session_state.video_bytes_cache = st.session_state.uploaded_video.getvalue()
-                st.session_state.uploaded_video.seek(0)
-            
-            with st.session_state.video_preview_placeholder.container():
+        with video_preview_container:
+            if st.session_state.uploaded_video is not None:
+                if st.session_state.video_bytes_cache is None:
+                    logger.info("📖 Caching video bytes...")
+                    st.session_state.video_bytes_cache = st.session_state.uploaded_video.getvalue()
+                    st.session_state.uploaded_video.seek(0)
+                    logger.info(f"✅ Cached {len(st.session_state.video_bytes_cache)/1e6:.2f}MB")
+                
                 st.video(st.session_state.video_bytes_cache)
-        else:
-            st.session_state.video_preview_placeholder.empty()
+            else:
+                st.info("No video uploaded yet")
     
     with col2:
         st.header("2. Background Settings")
@@ -264,26 +325,26 @@ def main():
             # Check if image actually changed using id()
             current_bg_id = id(bg_image)
             if current_bg_id != st.session_state.last_bg_image_id:
+                logger.info(f"🖼️ New background image: {bg_image.name if bg_image else 'None'}")
                 st.session_state.last_bg_image_id = current_bg_id
                 if bg_image is not None:
                     st.session_state.bg_image_cache = Image.open(bg_image)
                     st.session_state.bg_image_name = bg_image.name
+                    logger.info(f"✅ Background image cached: {bg_image.name}")
                 else:
                     st.session_state.bg_image_cache = None
             
-            # Background preview section
-            if st.session_state.bg_preview_placeholder is None:
-                st.session_state.bg_preview_placeholder = st.empty()
-            
-            if st.session_state.bg_image_cache is not None:
-                with st.session_state.bg_preview_placeholder.container():
+            # Background preview
+            bg_preview_container = st.container()
+            with bg_preview_container:
+                if st.session_state.bg_image_cache is not None:
                     st.image(
                         st.session_state.bg_image_cache,
                         caption="Selected Background",
                         use_container_width=True
                     )
-            else:
-                st.session_state.bg_preview_placeholder.empty()
+                else:
+                    st.info("No background image uploaded yet")
         
         elif bg_type == "Color":
             selected_color = st.color_picker(
@@ -294,6 +355,7 @@ def main():
             
             # Update only if color actually changed
             if selected_color != st.session_state.cached_color:
+                logger.info(f"🎨 Color changed to: {selected_color}")
                 st.session_state.bg_color = selected_color
                 st.session_state.cached_color = selected_color
                 
@@ -302,15 +364,11 @@ def main():
                 color_display[:, :] = color_rgb[::-1]
                 st.session_state.color_display_cache = color_display
             
-            # Color preview section
-            if st.session_state.bg_preview_placeholder is None:
-                st.session_state.bg_preview_placeholder = st.empty()
-            
-            if st.session_state.color_display_cache is not None:
-                with st.session_state.bg_preview_placeholder.container():
+            # Color preview
+            color_preview_container = st.container()
+            with color_preview_container:
+                if st.session_state.color_display_cache is not None:
                     st.image(st.session_state.color_display_cache, caption="Selected Color", width=200)
-            else:
-                st.session_state.bg_preview_placeholder.empty()
         
         st.header("3. Process & Download")
         
@@ -319,52 +377,69 @@ def main():
                 "🚀 Process Video",
                 disabled=not st.session_state.uploaded_video or st.session_state.processing
             )
-            process_status = st.empty()
             
             if submitted and not st.session_state.processing:
+                logger.info("🚀 PROCESS BUTTON CLICKED")
                 st.session_state.processing = True
+                st.session_state.process_complete = False
+                st.session_state.processed_video_bytes = None
+                
                 with st.spinner("Processing video (this may take a few minutes)..."):
                     try:
                         background = None
-                        if bg_type == "image" and st.session_state.bg_image_cache is not None:
+                        if bg_type == "Image" and st.session_state.bg_image_cache is not None:
                             background = st.session_state.bg_image_cache
-                        elif bg_type == "color" and 'bg_color' in st.session_state:
+                            logger.info("📦 Using background IMAGE")
+                        elif bg_type == "Color" and 'bg_color' in st.session_state:
                             background = st.session_state.bg_color
+                            logger.info(f"📦 Using background COLOR: {background}")
                         
-                        output_path = process_video(
+                        video_bytes = process_video(
                             st.session_state.uploaded_video,
                             background,
                             bg_type=bg_type.lower()
                         )
                         
-                        if output_path and os.path.exists(output_path):
-                            st.session_state.processed_video_path = output_path
-                            process_status.success("✅ Video processing complete!")
+                        if video_bytes and len(video_bytes) > 0:
+                            st.session_state.processed_video_bytes = video_bytes
+                            st.session_state.process_complete = True
+                            logger.info(f"✅ Processing complete! Video size: {len(video_bytes)/1e6:.2f}MB")
+                            st.success("✅ Video processing complete!")
                         else:
-                            process_status.error("❌ Failed to process video. Please check the logs for details.")
+                            logger.error("❌ Processing returned empty or None")
+                            st.error("❌ Failed to process video. Please check the logs for details.")
+                            
                     except Exception as e:
-                        process_status.error(f"❌ An error occurred: {str(e)}")
-                        logger.exception("Video processing failed")
+                        logger.error(f"❌ Exception during processing: {str(e)}")
+                        logger.error(traceback.format_exc())
+                        st.error(f"❌ An error occurred: {str(e)}")
                     finally:
                         st.session_state.processing = False
+                        logger.info(f"🏁 Processing finished. Success: {st.session_state.process_complete}")
         
-        if st.session_state.processed_video_path:
-            st.markdown("### Processed Video")
+        # Show processed video if available (OUTSIDE the form)
+        if st.session_state.processed_video_bytes is not None and len(st.session_state.processed_video_bytes) > 0:
+            st.markdown("---")
+            st.markdown("### ✅ Processed Video")
+            
             try:
-                with open(st.session_state.processed_video_path, 'rb') as f:
-                    video_bytes = f.read()
-                    st.video(video_bytes)
-                    st.download_button(
-                        label="💾 Download Processed Video",
-                        data=video_bytes,
-                        file_name="processed_video.mp4",
-                        mime="video/mp4",
-                        use_container_width=True,
-                        key="download_button"
-                    )
+                logger.info(f"📺 Displaying processed video: {len(st.session_state.processed_video_bytes)/1e6:.2f}MB")
+                st.video(st.session_state.processed_video_bytes)
+                
+                st.download_button(
+                    label="💾 Download Processed Video",
+                    data=st.session_state.processed_video_bytes,
+                    file_name="processed_video.mp4",
+                    mime="video/mp4",
+                    use_container_width=True,
+                    key="download_button"
+                )
+                logger.info("✅ Video displayed successfully")
+                
             except Exception as e:
+                logger.error(f"❌ Error displaying video: {str(e)}")
+                logger.error(traceback.format_exc())
                 st.error(f"Error displaying video: {str(e)}")
-                logger.error(f"Error displaying video: {str(e)}", exc_info=True)
 
 if __name__ == "__main__":
     main()