Update streamlit_app.py

streamlit_app.py

@@ -53,7 +53,6 @@ def check_gpu():
         logger.info(f"Device Name: {torch.cuda.get_device_name(0)}")
         logger.info(f"Device Capability: {torch.cuda.get_device_capability(0)}")
         
-        # Set as default device
         torch.cuda.set_device(0)
         logger.info("Set CUDA device 0 as default")
         
@@ -102,13 +101,17 @@ def initialize_session_state():
     defaults = {
         'uploaded_video': None,
         'video_bytes_cache': None,
+        'video_preview_placeholder': None,
         'bg_image_cache': None,
+        'bg_preview_placeholder': None,
         'bg_color': "#00FF00",
+        'cached_color': None,
         'color_display_cache': None,
         'processed_video_bytes': None,
         'processing': False,
         'gpu_available': None,
-        'files_loaded': False
+        'last_video_id': None,
+        'last_bg_image_id': None
     }
     for key, value in defaults.items():
         if key not in st.session_state:
@@ -125,12 +128,10 @@ def process_video(input_file, background, bg_type="image"):
     logger.info("=" * 60)
     
     try:
-        # Ensure GPU is default device
         if torch.cuda.is_available():
             torch.cuda.set_device(0)
             logger.info("GPU set as default device for processing")
         
-        # Create temp directory
         temp_base = Path(tempfile.gettempdir()) / "video_processing"
         temp_base.mkdir(exist_ok=True)
         temp_dir = temp_base / f"session_{int(time.time())}"
@@ -138,7 +139,6 @@ def process_video(input_file, background, bg_type="image"):
         
         logger.info(f"Temp directory: {temp_dir}")
 
-        # Save video
         input_path = str(temp_dir / "input.mp4")
         logger.info(f"Saving video to {input_path}")
         
@@ -150,7 +150,6 @@ def process_video(input_file, background, bg_type="image"):
         file_size = os.path.getsize(input_path)
         logger.info(f"Video saved: {file_size} bytes ({file_size/1e6:.2f}MB)")
 
-        # Prepare background
         bg_path = None
         if bg_type == "image" and background is not None:
             logger.info("Processing background IMAGE")
@@ -167,7 +166,6 @@ def process_video(input_file, background, bg_type="image"):
             cv2.imwrite(bg_path, np.ones((100, 100, 3), dtype=np.uint8) * color_rgb[::-1])
             logger.info(f"Background color saved")
 
-        # Progress tracking
         progress_bar = st.progress(0)
         status_text = st.empty()
         
@@ -177,18 +175,15 @@ def progress_callback(progress, message):
             progress_bar.progress(progress)
             status_text.text(f"{message} ({progress*100:.0f}%)")
 
-        # GPU status before processing
         if torch.cuda.is_available():
             logger.info(f"GPU Memory Before: {torch.cuda.memory_allocated()/1e9:.2f}GB")
 
-        # Process video
         output_path = str(temp_dir / "output.mp4")
         click_points = [[0.5, 0.5]]
         
         logger.info("Importing TwoStageProcessor...")
         from pipeline.integrated_pipeline import TwoStageProcessor
 
-        # Don't cache processor - causes GPU issues
         logger.info("Creating processor instance...")
         processor = TwoStageProcessor(temp_dir=str(temp_dir))
         logger.info("Processor created")
@@ -216,21 +211,18 @@ def progress_callback(progress, message):
         if not success:
             raise RuntimeError("Processing returned False")
 
-        # Check output
         if not os.path.exists(output_path):
             raise FileNotFoundError("Output video not created")
         
         output_size = os.path.getsize(output_path)
         logger.info(f"Output: {output_size} bytes ({output_size/1e6:.2f}MB)")
 
-        # Read output
         logger.info("Reading output...")
         with open(output_path, 'rb') as f:
             video_bytes = f.read()
         
         logger.info(f"Output loaded: {len(video_bytes)/1e6:.2f}MB")
 
-        # Cleanup
         try:
             shutil.rmtree(temp_dir)
             logger.info("Temp cleaned")
@@ -259,7 +251,6 @@ def main():
     
     initialize_session_state()
     
-    # GPU Status
     with st.sidebar:
         st.subheader("System Status")
         if st.session_state.gpu_available:
@@ -280,20 +271,20 @@ def main():
             key="video_uploader"
         )
         
-        # Handle file upload
-        if uploaded is not None:
-            # Always refresh if file exists (handles Streamlit restore bug)
-            if (st.session_state.uploaded_video is None or 
-                not hasattr(st.session_state.uploaded_video, 'name') or
-                st.session_state.uploaded_video.name != uploaded.name):
-                logger.info(f"Loading 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.files_loaded = True
+        # 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.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
         
-        # Video preview
+        # Video preview with placeholder
         st.markdown("### Video Preview")
+        if st.session_state.video_preview_placeholder is None:
+            st.session_state.video_preview_placeholder = st.empty()
+        
         if st.session_state.uploaded_video is not None:
             try:
                 if st.session_state.video_bytes_cache is None:
@@ -302,12 +293,13 @@ def main():
                     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)
+                with st.session_state.video_preview_placeholder.container():
+                    st.video(st.session_state.video_bytes_cache)
             except Exception as e:
                 logger.error(f"Video preview error: {e}")
-                st.error(f"Cannot display video: {e}")
+                st.session_state.video_preview_placeholder.error(f"Cannot display video: {e}")
         else:
-            st.info("No video uploaded")
+            st.session_state.video_preview_placeholder.empty()
     
     with col2:
         st.header("2. Background Settings")
@@ -326,18 +318,25 @@ def main():
                 key="bg_image_uploader"
             )
             
-            if bg_image is not None:
-                if (st.session_state.bg_image_cache is None or
-                    not hasattr(st.session_state.bg_image_cache, 'filename') or
-                    getattr(st.session_state.bg_image_cache, 'filename', None) != bg_image.name):
-                    logger.info(f"Loading background: {bg_image.name}")
-                    img = Image.open(bg_image)
-                    img.filename = bg_image.name
-                    st.session_state.bg_image_cache = img
-                
-                st.image(st.session_state.bg_image_cache, caption="Selected Background", use_container_width=True)
+            # 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: {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)
+                else:
+                    st.session_state.bg_image_cache = None
+            
+            # Background preview with placeholder
+            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():
+                    st.image(st.session_state.bg_image_cache, caption="Selected Background", use_container_width=True)
             else:
-                st.info("No background image")
+                st.session_state.bg_preview_placeholder.empty()
         
         elif bg_type == "Color":
             selected_color = st.color_picker(
@@ -346,15 +345,24 @@ def main():
                 key="color_picker"
             )
             
-            if selected_color != st.session_state.bg_color:
+            if selected_color != st.session_state.cached_color:
+                logger.info(f"Color: {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
             
+            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:
-                st.image(st.session_state.color_display_cache, caption="Selected Color", width=200)
+                with st.session_state.bg_preview_placeholder.container():
+                    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")
         
@@ -400,7 +408,6 @@ def main():
                 finally:
                     st.session_state.processing = False
         
-        # Show processed video
         if st.session_state.processed_video_bytes is not None:
             st.markdown("---")
             st.markdown("### Processed Video")
@@ -419,4 +426,4 @@ def main():
                 st.error(f"Display error: {e}")
 
 if __name__ == "__main__":
-    main() 
+    main()