Update ui/ui_components.py

ui/ui_components.py

@@ -22,6 +22,7 @@
     cb_use_gen_bg,
     cb_video_changed,
     cb_custom_bg_preview,
+    cb_preset_bg_preview,  # Added import for preset preview
 )
 
 CSS = """
@@ -60,9 +61,6 @@ def create_interface() -> gr.Blocks:
         analytics_enabled=False,
         theme=gr.themes.Soft()
     ) as demo:
-        
-        # State to hold the current background preview
-        background_state = gr.State()
 
         # ------------------------------------------------------------------
         # HERO
@@ -105,6 +103,12 @@ def create_interface() -> gr.Blocks:
                         interactive=True
                     )
 
+                    custom_bg_preview_img = gr.Image(
+                        label="Custom Background Preview",
+                        interactive=False,
+                        elem_classes=["preview-img"]
+                    )
+
                     # ── Advanced options accordion ───────────────────────
                     with gr.Accordion("Advanced", open=False):
                         use_two_stage = gr.Checkbox(
@@ -142,15 +146,6 @@ def create_interface() -> gr.Blocks:
 
                 # ── Right column ──────────────────────────────────────────
                 with gr.Column(scale=1):
-                    # Persistent background preview at the top
-                    gr.Markdown("### 📌 Background Preview")
-                    persistent_bg_preview = gr.Image(
-                        label="Selected Background (stays visible during processing)",
-                        interactive=False,
-                        elem_classes=["preview-img", "bg-preview-persistent"],
-                        show_label=True
-                    )
-                    
                     out_video = gr.Video(label="Processed Output", interactive=False)
                     statusbox = gr.Textbox(label="Status", lines=8, elem_id="statusbox")
 
@@ -209,36 +204,9 @@ def create_interface() -> gr.Blocks:
         # ------------------------------------------------------------------
         # Callback wiring
         # ------------------------------------------------------------------
-        
-        # Modified callback to update persistent preview when custom background changes
-        def update_persistent_preview(custom_bg_file):
-            """Update both the temporary preview and persistent preview"""
-            if custom_bg_file is None:
-                return None, None, None
-            try:
-                import cv2
-                img = cv2.imread(custom_bg_file.name)
-                if img is not None:
-                    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-                    return img_rgb, img_rgb, custom_bg_file  # Return for both previews and state
-            except:
-                pass
-            return None, None, None
-        
-        # Update persistent preview when background style changes (for presets)
-        def update_preset_preview(bg_style_choice):
-            """Generate preview for preset backgrounds"""
-            try:
-                from utils.cv_processing import create_professional_background
-                # Create a small preview version
-                preview_bg = create_professional_background(bg_style_choice, 640, 360)
-                return preview_bg
-            except:
-                return None
-        
         btn_load.click(cb_load_models, outputs=statusbox)
 
-        # Modified process callback to preserve background preview
+        # CRITICAL FIX: Process button now expects THREE outputs from callback
         btn_run.click(
             cb_process_video,
             inputs=[
@@ -251,16 +219,16 @@ def update_preset_preview(bg_style_choice):
                 preview_mask,
                 preview_greenscreen,
             ],
-            outputs=[out_video, statusbox],  # Note: NOT clearing persistent_bg_preview
+            outputs=[out_video, statusbox, custom_bg_preview_img],  # THREE outputs now!
         )
 
         btn_cancel.click(cb_cancel, outputs=statusbox)
         btn_refresh.click(cb_status, outputs=[model_status, cache_status])
         
-        # Modified clear to also clear the persistent preview
+        # Clear button - adjusted for 5 outputs from callback
         btn_clear.click(
             cb_clear, 
-            outputs=[out_video, statusbox, gen_preview, gen_path, persistent_bg_preview]
+            outputs=[out_video, statusbox, gen_preview, gen_path, custom_bg_preview_img]
         )
 
         # AI background generation
@@ -270,54 +238,24 @@ def update_preset_preview(bg_style_choice):
             outputs=[gen_preview, gen_path],
         )
         
-        # Modified to update persistent preview when using generated background
-        def use_generated_as_custom(gen_path_value):
-            """Use generated background as custom and update persistent preview"""
-            if gen_path_value:
-                try:
-                    import cv2
-                    img = cv2.imread(gen_path_value)
-                    if img is not None:
-                        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-                        # Create a file-like object for the custom_bg component
-                        import tempfile
-                        import shutil
-                        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
-                        cv2.imwrite(temp_file.name, img)
-                        temp_file.close()
-                        return temp_file.name, img_rgb
-                except:
-                    pass
-            return None, None
-        
-        use_gen_as_custom.click(
-            use_generated_as_custom, 
-            inputs=[gen_path], 
-            outputs=[custom_bg, persistent_bg_preview]
-        )
+        use_gen_as_custom.click(cb_use_gen_bg, inputs=[gen_path], outputs=[custom_bg])
 
         # Live previews
         video.change(cb_video_changed, inputs=[video], outputs=[video_preview])
+        custom_bg.change(cb_custom_bg_preview, inputs=[custom_bg], outputs=[custom_bg_preview_img])
         
-        # Update persistent preview when custom background changes
-        custom_bg.change(
-            update_persistent_preview, 
-            inputs=[custom_bg], 
-            outputs=[video_preview, persistent_bg_preview, background_state]
-        )
-        
-        # Update persistent preview when background style dropdown changes
+        # Update preview when background style dropdown changes
         bg_style.change(
-            update_preset_preview,
+            cb_preset_bg_preview,
             inputs=[bg_style],
-            outputs=[persistent_bg_preview]
+            outputs=[custom_bg_preview_img]
         )
         
-        # Initialize with default background preview
+        # Initialize with default background preview on load
         demo.load(
-            update_preset_preview,
+            cb_preset_bg_preview,
             inputs=[bg_style],
-            outputs=[persistent_bg_preview]
+            outputs=[custom_bg_preview_img]
         )
 
     return demo