Update streamlit_app.py

streamlit_app.py

@@ -30,12 +30,45 @@ def custom_excepthook(type, value, tb):
 
 # --- Streamlit Page Config ---
 st.set_page_config(
-    page_title="🎬 Advanced Video Background Replacer",
+    page_title="Advanced Video Background Replacer",
     page_icon="🎥",
     layout="wide",
     initial_sidebar_state="expanded"
 )
 
+# --- GPU Diagnostic ---
+def check_gpu():
+    """Check GPU availability and log details"""
+    logger.info("=" * 60)
+    logger.info("GPU DIAGNOSTIC")
+    logger.info("=" * 60)
+    
+    cuda_available = torch.cuda.is_available()
+    logger.info(f"torch.cuda.is_available(): {cuda_available}")
+    
+    if cuda_available:
+        logger.info(f"CUDA Version: {torch.version.cuda}")
+        logger.info(f"Device Count: {torch.cuda.device_count()}")
+        logger.info(f"Current Device: {torch.cuda.current_device()}")
+        logger.info(f"Device Name: {torch.cuda.get_device_name(0)}")
+        logger.info(f"Device Capability: {torch.cuda.get_device_capability(0)}")
+        
+        # Test GPU tensor
+        try:
+            test_tensor = torch.randn(100, 100).cuda()
+            logger.info(f"GPU Test Tensor Created Successfully on: {test_tensor.device}")
+            del test_tensor
+            torch.cuda.empty_cache()
+        except Exception as e:
+            logger.error(f"GPU Test Failed: {e}")
+    else:
+        logger.warning("CUDA NOT AVAILABLE")
+        logger.info(f"PyTorch Version: {torch.__version__}")
+        logger.info(f"CUDA Built Version: {torch.version.cuda}")
+    
+    logger.info("=" * 60)
+    return cuda_available
+
 # --- Custom CSS ---
 st.markdown("""
 <style>
@@ -59,21 +92,6 @@ def custom_excepthook(type, value, tb):
     .stAlert {
         border-radius: 10px;
     }
-    .stTabs [data-baseweb="tab-list"] {
-        gap: 10px;
-    }
-    .stTabs [data-baseweb="tab"] {
-        height: 50px;
-        white-space: pre;
-        background-color: #f0f2f6;
-        border-radius: 4px 4px 0 0;
-        padding: 10px 20px;
-        margin-right: 5px;
-    }
-    .stTabs [aria-selected="true"] {
-        background-color: #4CAF50;
-        color: white;
-    }
     .video-container {
         border: 2px dashed #4CAF50;
         border-radius: 10px;
@@ -86,35 +104,26 @@ def custom_excepthook(type, value, tb):
 # --- Session State Initialization ---
 def initialize_session_state():
     """Initialize all session state variables"""
-    defaults = {
-        'uploaded_video': None,
-        'video_bytes_cache': None,
-        'bg_image': None,
-        'bg_image_cache': None,
-        'bg_image_name': None,
-        'bg_color': "#00FF00",
-        'cached_color': None,
-        'color_display_cache': 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,
-        'process_complete': False
-    }
-    for key, value in defaults.items():
-        if key not in st.session_state:
-            st.session_state[key] = value
+    if 'initialized' not in st.session_state:
+        st.session_state.initialized = True
+        st.session_state.uploaded_video = None
+        st.session_state.video_bytes_cache = None
+        st.session_state.bg_image_cache = None
+        st.session_state.bg_color = "#00FF00"
+        st.session_state.color_display_cache = None
+        st.session_state.processed_video_bytes = None
+        st.session_state.processing = False
+        st.session_state.process_complete = False
+        
+        # Run GPU check on first init
+        gpu_available = check_gpu()
+        st.session_state.gpu_available = gpu_available
 
 # --- Video Processing ---
 def process_video(input_file, background, bg_type="image"):
-    """
-    Process video with the selected background using SAM2 and MatAnyone pipeline.
-    Returns the video bytes.
-    """
+    """Process video with the selected background using SAM2 and MatAnyone pipeline."""
     logger.info("=" * 60)
-    logger.info("🎬 STARTING VIDEO PROCESSING")
+    logger.info("STARTING VIDEO PROCESSING")
     logger.info("=" * 60)
     
     try:
@@ -124,77 +133,78 @@ def process_video(input_file, background, bg_type="image"):
         temp_dir = temp_base / f"session_{int(time.time())}"
         temp_dir.mkdir(exist_ok=True)
         
-        logger.info(f"📁 Created temp directory: {temp_dir}")
+        logger.info(f"Temp directory: {temp_dir}")
 
-        # Save the uploaded video to a temporary file
+        # Save the uploaded video
         input_path = str(temp_dir / "input.mp4")
-        logger.info(f"💾 Writing video to {input_path}")
+        logger.info(f"Writing video to {input_path}")
         
+        input_file.seek(0)  # Ensure we're at the start
         with open(input_path, "wb") as f:
-            written = f.write(input_file.getvalue())
-            logger.info(f"✅ Wrote {written/1e6:.2f}MB")
+            written = f.write(input_file.read())
+            logger.info(f"Wrote {written/1e6:.2f}MB")
         
-        if not os.path.exists(input_path):
-            raise FileNotFoundError(f"❌ Input video not saved: {input_path}")
+        if not os.path.exists(input_path) or os.path.getsize(input_path) == 0:
+            raise FileNotFoundError(f"Input video not saved properly: {input_path}")
         
-        logger.info(f"✅ Video file exists: {os.path.getsize(input_path)} bytes")
+        logger.info(f"Video file verified: {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")
+            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}")
+            logger.info(f"Background image written: {bg_path}")
             
-        elif bg_type == "color" and hasattr(st.session_state, 'bg_color'):
-            logger.info(f"🎨 Processing background COLOR: {st.session_state.bg_color}")
+        elif bg_type == "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}")
+            logger.info(f"Background color image written: {bg_path}")
 
-        # Set up progress placeholders
+        # Progress tracking
         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}")
+            logger.info(f"Progress: {progress*100:.1f}% - {message}")
             progress_placeholder.progress(progress)
             status_placeholder.text(f"Status: {message}")
 
-        # Log GPU state
+        # GPU check
         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")
+            device = torch.cuda.get_device_name(0)
+            mem_before = torch.cuda.memory_allocated()/1e9
+            logger.info(f"CUDA Device: {device}")
+            logger.info(f"GPU Memory Before: {mem_before:.2f}GB")
         else:
-            logger.info("⚠️ CUDA Available: False (using CPU)")
+            logger.warning("CUDA NOT AVAILABLE - Processing on CPU")
 
         # Process the video
         output_path = str(temp_dir / "output.mp4")
         click_points = [[0.5, 0.5]]
         
-        logger.info("🔧 Importing TwoStageProcessor...")
+        logger.info("Importing TwoStageProcessor...")
         from pipeline.integrated_pipeline import TwoStageProcessor
 
-        # Cache processor
         @st.cache_resource
         def load_processor(temp_dir_str):
-            logger.info(f"🏗️ Loading processor with temp_dir: {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")
+        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}")
+            logger.info("Calling process_video...")
+            logger.info(f"  Input: {input_path}")
+            logger.info(f"  Background: {bg_path}")
+            logger.info(f"  Output: {output_path}")
             
             success = processor.process_video(
                 input_video=input_path,
@@ -205,52 +215,51 @@ def load_processor(temp_dir_str):
                 progress_callback=progress_callback
             )
             
-            logger.info(f"📊 Processing returned: {success}")
+            logger.info(f"Processing returned: {success}")
             
         except Exception as e:
-            logger.error(f"❌ Pipeline processing failed: {traceback.format_exc()}", exc_info=True)
+            logger.error(f"Pipeline processing failed: {traceback.format_exc()}")
             raise
 
         if torch.cuda.is_available():
-            logger.info(f"💾 GPU Mem After: {torch.cuda.memory_allocated()/1e9:.2f}GB")
+            mem_after = torch.cuda.memory_allocated()/1e9
+            logger.info(f"GPU Memory After: {mem_after:.2f}GB")
             torch.cuda.empty_cache()
-            logger.info("🧹 GPU cache cleared")
 
         if not success:
-            raise RuntimeError("❌ Video processing returned False")
+            raise RuntimeError("Video processing returned False")
 
-        # Check if output file exists
+        # Verify output
         if not os.path.exists(output_path):
-            logger.error(f"❌ Output file does not exist: {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)")
+        logger.info(f"Output file exists: {output_size} bytes ({output_size/1e6:.2f}MB)")
 
-        # Read the output video into memory
-        logger.info("📖 Reading output video into memory...")
+        # Read output 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")
+        logger.info(f"Read {len(video_bytes)/1e6:.2f}MB into memory")
 
-        # Clean up temp directory
+        # Cleanup
         try:
-            logger.info(f"🧹 Cleaning up temp directory: {temp_dir}")
             shutil.rmtree(temp_dir)
-            logger.info("✅ Temp directory cleaned")
+            logger.info("Temp directory cleaned")
         except Exception as e:
-            logger.warning(f"⚠️ Could not clean temp directory: {e}")
+            logger.warning(f"Could not clean temp directory: {e}")
 
         logger.info("=" * 60)
-        logger.info("✅ VIDEO PROCESSING COMPLETED SUCCESSFULLY")
+        logger.info("VIDEO PROCESSING COMPLETED")
         logger.info("=" * 60)
         
         return video_bytes
 
     except Exception as e:
         logger.error("=" * 60)
-        logger.error(f"❌ ERROR IN VIDEO PROCESSING: {str(e)}")
+        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)}")
@@ -258,13 +267,21 @@ def load_processor(temp_dir_str):
 
 # --- Main Application ---
 def main():
-    st.title("🎬 Advanced Video Background Replacer")
+    st.title("Advanced Video Background Replacer")
     st.markdown("---")
     
     # Initialize session state
     initialize_session_state()
     
-    logger.info(f"🔄 App rerun - Processing: {st.session_state.processing}, Complete: {st.session_state.process_complete}")
+    # GPU Status in sidebar
+    with st.sidebar:
+        st.subheader("System Status")
+        if st.session_state.gpu_available:
+            st.success(f"GPU: {torch.cuda.get_device_name(0)}")
+        else:
+            st.error("GPU: Not Available (using CPU)")
+    
+    logger.info(f"Rerun - Processing: {st.session_state.processing}, Complete: {st.session_state.process_complete}")
     
     # Main layout
     col1, col2 = st.columns([1, 1], gap="large")
@@ -273,36 +290,32 @@ def main():
         st.header("1. Upload Video")
         
         uploaded = st.file_uploader(
-            "📹 Upload Video",
+            "Upload Video",
             type=["mp4", "mov", "avi"],
             key="video_uploader"
         )
         
-        # 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
+        # Store uploaded video
+        if uploaded is not None:
+            if st.session_state.uploaded_video is None or uploaded.name != getattr(st.session_state.uploaded_video, 'name', ''):
+                logger.info(f"New video: {uploaded.name} ({uploaded.size} bytes)")
+                st.session_state.uploaded_video = uploaded
+                st.session_state.video_bytes_cache = None
+                st.session_state.processed_video_bytes = None
+                st.session_state.process_complete = False
         
-        # Video preview section
+        # Video preview
         st.markdown("### Video Preview")
-        video_preview_container = st.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.info("No video uploaded yet")
+        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.uploaded_video.seek(0)
+                st.session_state.video_bytes_cache = st.session_state.uploaded_video.read()
+                logger.info(f"Cached {len(st.session_state.video_bytes_cache)/1e6:.2f}MB")
+            
+            st.video(st.session_state.video_bytes_cache)
+        else:
+            st.info("No video uploaded yet")
     
     with col2:
         st.header("2. Background Settings")
@@ -311,134 +324,110 @@ def main():
             "Select Background Type:",
             ["Image", "Color", "Blur"],
             horizontal=True,
-            index=0,
             key="bg_type_radio"
         )
         
         if bg_type == "Image":
             bg_image = st.file_uploader(
-                "🖼️ Upload Background Image",
+                "Upload Background Image",
                 type=["jpg", "png", "jpeg"],
                 key="bg_image_uploader"
             )
             
-            # 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:
+            if bg_image is not None:
+                if st.session_state.bg_image_cache is None or bg_image.name != getattr(st.session_state.bg_image_cache, 'name', ''):
+                    logger.info(f"New background: {bg_image.name}")
                     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
-            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.info("No background image uploaded yet")
+                    st.session_state.bg_image_cache.name = bg_image.name
+                
+                st.image(st.session_state.bg_image_cache, caption="Selected Background", use_container_width=True)
+            else:
+                st.info("No background image uploaded yet")
         
         elif bg_type == "Color":
             selected_color = st.color_picker(
-                "🎨 Choose Background Color",
+                "Choose Background Color",
                 st.session_state.bg_color,
                 key="color_picker"
             )
             
-            # Update only if color actually changed
-            if selected_color != st.session_state.cached_color:
-                logger.info(f"🎨 Color changed to: {selected_color}")
+            if selected_color != st.session_state.bg_color:
+                logger.info(f"Color changed to: {selected_color}")
                 st.session_state.bg_color = selected_color
-                st.session_state.cached_color = selected_color
                 
                 color_rgb = tuple(int(selected_color.lstrip('#')[i:i+2], 16) for i in (0, 2, 4))
                 color_display = np.zeros((100, 100, 3), dtype=np.uint8)
                 color_display[:, :] = color_rgb[::-1]
                 st.session_state.color_display_cache = color_display
             
-            # 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)
+            if st.session_state.color_display_cache is not None:
+                st.image(st.session_state.color_display_cache, caption="Selected Color", width=200)
         
         st.header("3. Process & Download")
         
-        with st.form("process_form"):
-            submitted = st.form_submit_button(
-                "🚀 Process Video",
-                disabled=not st.session_state.uploaded_video or st.session_state.processing
-            )
+        # Simple button instead of form
+        if st.button(
+            "Process Video",
+            disabled=st.session_state.uploaded_video is None or st.session_state.processing,
+            use_container_width=True
+        ):
+            logger.info("PROCESS BUTTON CLICKED")
+            st.session_state.processing = True
+            st.session_state.process_complete = False
+            st.session_state.processed_video_bytes = None
             
-            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:
-                            background = st.session_state.bg_image_cache
-                            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}")
-                        
-                        video_bytes = process_video(
-                            st.session_state.uploaded_video,
-                            background,
-                            bg_type=bg_type.lower()
-                        )
+            with st.spinner("Processing video..."):
+                try:
+                    background = None
+                    if bg_type == "Image" and st.session_state.bg_image_cache is not None:
+                        background = st.session_state.bg_image_cache
+                        logger.info("Using background IMAGE")
+                    elif bg_type == "Color":
+                        background = st.session_state.bg_color
+                        logger.info(f"Using background COLOR: {background}")
+                    
+                    video_bytes = process_video(
+                        st.session_state.uploaded_video,
+                        background,
+                        bg_type=bg_type.lower()
+                    )
+                    
+                    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! {len(video_bytes)/1e6:.2f}MB")
+                        st.success("Video processing complete!")
+                    else:
+                        logger.error("Processing returned empty or None")
+                        st.error("Failed to process video")
                         
-                        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:
-                            logger.error("❌ Processing returned empty or None")
-                            st.error("❌ Failed to process video. Please check the logs for details.")
-                            
-                    except Exception as e:
-                        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}")
+                except Exception as e:
+                    logger.error(f"Exception: {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}")
         
-        # Show processed video if available (OUTSIDE the form)
+        # Show processed video
         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")
+            st.markdown("### Processed Video")
             
             try:
-                logger.info(f"📺 Displaying processed video: {len(st.session_state.processed_video_bytes)/1e6:.2f}MB")
+                logger.info(f"Displaying 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",
+                    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"
+                    use_container_width=True
                 )
-                logger.info("✅ Video displayed successfully")
                 
             except Exception as e:
-                logger.error(f"❌ Error displaying video: {str(e)}")
-                logger.error(traceback.format_exc())
+                logger.error(f"Error displaying video: {str(e)}")
                 st.error(f"Error displaying video: {str(e)}")
 
 if __name__ == "__main__":