Update ui/ui_components.py

ui/ui_components.py

@@ -6,6 +6,7 @@
 * All callbacks live in ui/callbacks.py
 * Adds live previews for the first frame of the uploaded video
   and for the custom background image
+* UPDATED: Background preview persists during processing
 """
 
 from __future__ import annotations
@@ -34,6 +35,11 @@
 .sm { font-size: 13px; opacity: 0.85; }
 #statusbox { min-height: 120px; }
 .preview-img { border-radius: var(--radius); border: 1px solid rgba(0,0,0,.08); }
+.bg-preview-persistent { 
+    border: 2px solid rgba(0,120,255,0.3); 
+    background: rgba(0,120,255,0.05);
+    min-height: 200px;
+}
 """
 
 # Keep in sync with utils/cv_processing.PROFESSIONAL_BACKGROUNDS
@@ -54,6 +60,9 @@ 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
@@ -96,12 +105,6 @@ 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(
@@ -139,6 +142,15 @@ 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")
 
@@ -197,8 +209,36 @@ 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
         btn_run.click(
             cb_process_video,
             inputs=[
@@ -207,16 +247,21 @@ def create_interface() -> gr.Blocks:
                 custom_bg,
                 use_two_stage,
                 chroma_preset,
-                key_color_mode,          # NEW
+                key_color_mode,
                 preview_mask,
                 preview_greenscreen,
             ],
-            outputs=[out_video, statusbox],
+            outputs=[out_video, statusbox],  # Note: NOT clearing persistent_bg_preview
         )
 
         btn_cancel.click(cb_cancel, outputs=statusbox)
         btn_refresh.click(cb_status, outputs=[model_status, cache_status])
-        btn_clear.click(cb_clear, outputs=[out_video, statusbox, gen_preview, gen_path])
+        
+        # Modified clear to also clear the persistent preview
+        btn_clear.click(
+            cb_clear, 
+            outputs=[out_video, statusbox, gen_preview, gen_path, persistent_bg_preview]
+        )
 
         # AI background generation
         btn_gen_bg.click(
@@ -224,10 +269,55 @@ def create_interface() -> gr.Blocks:
             inputs=[prompt, gen_width, gen_height, bokeh, vignette, contrast],
             outputs=[gen_preview, gen_path],
         )
-        use_gen_as_custom.click(cb_use_gen_bg, inputs=[gen_path], outputs=[custom_bg])
+        
+        # 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]
+        )
 
         # 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
+        bg_style.change(
+            update_preset_preview,
+            inputs=[bg_style],
+            outputs=[persistent_bg_preview]
+        )
+        
+        # Initialize with default background preview
+        demo.load(
+            update_preset_preview,
+            inputs=[bg_style],
+            outputs=[persistent_bg_preview]
+        )
 
-    return demo
+    return demo