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)}")
         
-        # Test GPU tensor
         try:
             test_tensor = torch.randn(100, 100).cuda()
             logger.info(f"GPU Test Tensor Created Successfully on: {test_tensor.device}")
@@ -104,20 +103,29 @@ def check_gpu():
 # --- Session State Initialization ---
 def initialize_session_state():
     """Initialize all session state variables"""
-    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
+    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,
+        'process_complete': False,
+        'last_video_id': None,
+        'last_bg_image_id': None,
+        'gpu_available': None
+    }
+    for key, value in defaults.items():
+        if key not in st.session_state:
+            st.session_state[key] = value
+    
+    # Run GPU check on first init
+    if st.session_state.gpu_available is None:
+        st.session_state.gpu_available = check_gpu()
 
 # --- Video Processing ---
 def process_video(input_file, background, bg_type="image"):
@@ -139,13 +147,13 @@ def process_video(input_file, background, bg_type="image"):
         input_path = str(temp_dir / "input.mp4")
         logger.info(f"Writing video to {input_path}")
         
-        input_file.seek(0)  # Ensure we're at the start
+        input_file.seek(0)
         with open(input_path, "wb") as f:
             written = f.write(input_file.read())
             logger.info(f"Wrote {written/1e6:.2f}MB")
         
         if not os.path.exists(input_path) or os.path.getsize(input_path) == 0:
-            raise FileNotFoundError(f"Input video not saved properly: {input_path}")
+            raise FileNotFoundError(f"Input video not saved properly")
         
         logger.info(f"Video file verified: {os.path.getsize(input_path)} bytes")
 
@@ -156,7 +164,7 @@ def process_video(input_file, background, bg_type="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: {bg_path}")
+            logger.info(f"Background image written")
             
         elif bg_type == "color":
             logger.info(f"Processing background COLOR: {st.session_state.bg_color}")
@@ -164,7 +172,7 @@ def process_video(input_file, background, bg_type="image"):
             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: {bg_path}")
+            logger.info(f"Background color image written")
 
         # Progress tracking
         progress_placeholder = st.empty()
@@ -194,7 +202,7 @@ def progress_callback(progress, message):
 
         @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")
             return TwoStageProcessor(temp_dir=temp_dir_str)
         
         processor = load_processor(str(temp_dir))
@@ -202,9 +210,6 @@ def load_processor(temp_dir_str):
 
         try:
             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,
@@ -231,25 +236,24 @@ def load_processor(temp_dir_str):
 
         # Verify output
         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}")
+            raise FileNotFoundError(f"Output video not created")
         
         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: {output_size} bytes ({output_size/1e6:.2f}MB)")
 
         # Read output into memory
-        logger.info("Reading output video into memory...")
+        logger.info("Reading output 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")
 
         # Cleanup
         try:
             shutil.rmtree(temp_dir)
             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: {e}")
 
         logger.info("=" * 60)
         logger.info("VIDEO PROCESSING COMPLETED")
@@ -258,11 +262,9 @@ def load_processor(temp_dir_str):
         return video_bytes
 
     except Exception as e:
-        logger.error("=" * 60)
-        logger.error(f"ERROR IN VIDEO PROCESSING: {str(e)}")
+        logger.error(f"ERROR: {str(e)}")
         logger.error(traceback.format_exc())
-        logger.error("=" * 60)
-        st.error(f"An error occurred during processing: {str(e)}")
+        st.error(f"An error occurred: {str(e)}")
         return None
 
 # --- Main Application ---
@@ -279,9 +281,9 @@ def main():
         if st.session_state.gpu_available:
             st.success(f"GPU: {torch.cuda.get_device_name(0)}")
         else:
-            st.error("GPU: Not Available (using CPU)")
+            st.error("GPU: Not Available")
     
-    logger.info(f"Rerun - Processing: {st.session_state.processing}, Complete: {st.session_state.process_complete}")
+    logger.info(f"Rerun - Processing: {st.session_state.processing}")
     
     # Main layout
     col1, col2 = st.columns([1, 1], gap="large")
@@ -295,27 +297,32 @@ def main():
             key="video_uploader"
         )
         
-        # 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
+        # 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
+            st.session_state.process_complete = False
         
-        # 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:
             if st.session_state.video_bytes_cache is None:
-                logger.info("Caching video bytes...")
+                logger.info("Caching video...")
                 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)
+            with st.session_state.video_preview_placeholder.container():
+                st.video(st.session_state.video_bytes_cache)
         else:
-            st.info("No video uploaded yet")
+            st.session_state.video_preview_placeholder.empty()
     
     with col2:
         st.header("2. Background Settings")
@@ -334,15 +341,25 @@ def main():
                 key="bg_image_uploader"
             )
             
-            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}")
+            # 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)
-                    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.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 uploaded yet")
+                st.session_state.bg_preview_placeholder.empty()
         
         elif bg_type == "Color":
             selected_color = st.color_picker(
@@ -351,21 +368,27 @@ def main():
                 key="color_picker"
             )
             
-            if selected_color != st.session_state.bg_color:
-                logger.info(f"Color changed to: {selected_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")
         
-        # Simple button instead of form
         if st.button(
             "Process Video",
             disabled=st.session_state.uploaded_video is None or st.session_state.processing,
@@ -381,10 +404,8 @@ def main():
                     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,
@@ -395,19 +416,17 @@ def main():
                     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")
+                        logger.info(f"Complete! {len(video_bytes)/1e6:.2f}MB")
                         st.success("Video processing complete!")
                     else:
-                        logger.error("Processing returned empty or None")
+                        logger.error("Processing returned empty")
                         st.error("Failed to process video")
                         
                 except Exception as e:
                     logger.error(f"Exception: {str(e)}")
-                    logger.error(traceback.format_exc())
-                    st.error(f"An error occurred: {str(e)}")
+                    st.error(f"Error: {str(e)}")
                 finally:
                     st.session_state.processing = False
-                    logger.info(f"Processing finished. Success: {st.session_state.process_complete}")
         
         # Show processed video
         if st.session_state.processed_video_bytes is not None and len(st.session_state.processed_video_bytes) > 0:
@@ -415,9 +434,7 @@ def main():
             st.markdown("### Processed Video")
             
             try:
-                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",
                     data=st.session_state.processed_video_bytes,
@@ -425,10 +442,9 @@ def main():
                     mime="video/mp4",
                     use_container_width=True
                 )
-                
             except Exception as e:
-                logger.error(f"Error displaying video: {str(e)}")
-                st.error(f"Error displaying video: {str(e)}")
+                logger.error(f"Display error: {str(e)}")
+                st.error(f"Error: {str(e)}")
 
 if __name__ == "__main__":
     main()