Update ui.py

ui.py

@@ -1,10 +1,10 @@
 #!/usr/bin/env python3
 """
-Modern UI for Video Background Replacer
-- Clean, responsive layout with Streamlit
+Modern UI for Video Background Replacer (PRO)
+- Clean, responsive Streamlit layout
 - File uploaders, previews, and processing controls
-- Log viewer and download buttons
-- Passes proper PIL.Image or color string into pipeline
+- Supports: Color, Image upload, Stock image, AI Prompt background
+- Log viewer, session state, robust error handling
 """
 import streamlit as st
 import os
@@ -16,7 +16,6 @@
 
 logger = logging.getLogger("Advanced Video Background Replacer")
 
-# --- UI Helpers ---
 def tail_file(path: str, lines: int = 400) -> str:
     if not os.path.exists(path):
         return "(log file not found)"
@@ -36,6 +35,75 @@ def read_file_bytes(path: str) -> bytes:
     except Exception:
         return b""
 
+def _render_background_settings():
+    # Available stock images (replace with your real logic or API)
+    stock_images = {
+        "Sunset Beach": "stock_images/sunset_beach.jpg",
+        "Urban Office": "stock_images/urban_office.jpg",
+        "Studio Lighting": "stock_images/studio_light.jpg",
+    }
+    st.header("2. Background Settings")
+    bg_type = st.radio(
+        "Select Background Type:",
+        ["Image", "Color", "Stock", "AI Prompt"],
+        horizontal=True,
+        key="bg_type_radio"
+    )
+    background = None
+
+    if bg_type == "Image":
+        bg_image = st.file_uploader(
+            "Upload Background Image",
+            type=["jpg", "png", "jpeg"],
+            key="bg_image_uploader"
+        )
+        if bg_image is not None:
+            bg_image.seek(0)
+            background = Image.open(bg_image).convert("RGB")
+            st.image(background, caption="Selected Background", use_column_width=True)
+
+    elif bg_type == "Color":
+        selected_color = st.color_picker(
+            "Choose Background Color",
+            st.session_state.get('bg_color', "#00FF00"),
+            key="color_picker"
+        )
+        background = selected_color
+        color_preview = np.full(
+            (100, 100, 3),
+            tuple(int(selected_color.lstrip('#')[i:i+2], 16) for i in (0, 2, 4)),
+            dtype=np.uint8
+        )
+        st.image(color_preview, caption="Selected Color", width=200)
+
+    elif bg_type == "Stock":
+        stock_choice = st.selectbox(
+            "Choose a professional stock background:",
+            list(stock_images.keys()),
+            key="stock_image_select"
+        )
+        stock_img_path = stock_images[stock_choice]
+        background = Image.open(stock_img_path).convert("RGB")
+        st.image(background, caption=stock_choice, use_column_width=True)
+
+    elif bg_type == "AI Prompt":
+        prompt = st.text_input("Describe the background to generate (AI):", key="ai_bg_prompt")
+        ai_ready = False
+        if st.button("Generate Background", key="gen_bg_btn") and prompt:
+            # TODO: Plug in real AI image generator
+            # For now, stub with a purple placeholder
+            background = Image.new("RGB", (512, 320), (64, 32, 96))
+            st.session_state.generated_bg = background
+            st.success("AI-generated background (stub). Replace with your generator!")
+            ai_ready = True
+        elif "generated_bg" in st.session_state:
+            background = st.session_state.generated_bg
+            ai_ready = True
+        if ai_ready and background is not None:
+            st.image(background, caption="Generated Background", use_column_width=True)
+
+    return background, bg_type
+
 def render_ui(process_video_func):
     try:
         # --- Sidebar: System Status & Logs ---
@@ -91,59 +159,10 @@ def render_ui(process_video_func):
                     video_preview_placeholder.error(f"Cannot display video: {e}")
             else:
                 video_preview_placeholder.empty()
+
         # --- Column 2: Background & Processing ---
         with col2:
-            st.header("2. Background Settings")
-            bg_type = st.radio(
-                "Select Background Type:",
-                ["Image", "Color"],
-                horizontal=True,
-                key="bg_type_radio"
-            )
-            background = None
-            # --- Background Image Upload ---
-            if bg_type == "Image":
-                bg_image = st.file_uploader(
-                    "Upload Background Image",
-                    type=["jpg", "png", "jpeg"],
-                    key="bg_image_uploader"
-                )
-                bg_preview_placeholder = st.empty()
-                background_pil = None
-                if bg_image is not None:
-                    try:
-                        bg_image.seek(0)
-                        background_pil = Image.open(bg_image).convert("RGB")
-                        with bg_preview_placeholder.container():
-                            st.image(background_pil, caption="Selected Background", use_column_width=True)
-                    except Exception as e:
-                        logger.error(f"[UI] Background image error: {e}", exc_info=True)
-                        bg_preview_placeholder.error(f"Cannot display image: {e}")
-                        background_pil = None
-                else:
-                    bg_preview_placeholder.empty()
-                background = background_pil
-            else:
-                selected_color = st.color_picker(
-                    "Choose Background Color",
-                    st.session_state.get('bg_color', "#00FF00"),
-                    key="color_picker"
-                )
-                if selected_color != st.session_state.get('cached_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
-                    st.session_state.color_display_cache = color_display
-                color_preview_placeholder = st.empty()
-                if st.session_state.get('color_display_cache') is not None:
-                    with color_preview_placeholder.container():
-                        st.image(st.session_state.color_display_cache, caption="Selected Color", width=200)
-                else:
-                    color_preview_placeholder.empty()
-                background = selected_color  # "#RRGGBB"
-            # --- Process Button ---
+            background, bg_type = _render_background_settings()
             st.header("3. Process Video")
             can_process = (
                 uploaded_video is not None
@@ -157,7 +176,7 @@ def render_ui(process_video_func):
                         st.error("Please upload a background image")
                         return
                     with st.spinner("Processing video..."):
-                        # === FIX: Always pass lowercased bg_type as third argument ===
+                        # bg_type always lowercase for pipeline
                         success = process_video_func(uploaded_video, background, bg_type.lower())
                     if success:
                         st.success("Video processing complete!")
@@ -166,6 +185,7 @@ def render_ui(process_video_func):
                 except Exception as e:
                     logger.error(f"[UI] Process video error: {e}", exc_info=True)
                     st.error(f"Processing error: {str(e)}. Check logs for details.")
+
         # --- Results Display ---
         if st.session_state.get('processed_video_bytes') is not None:
             st.markdown("---")
@@ -182,6 +202,7 @@ def render_ui(process_video_func):
             except Exception as e:
                 logger.error(f"[UI] Display error: {e}", exc_info=True)
                 st.error(f"Display error: {e}")
+
         # --- Error Display ---
         if st.session_state.get('last_error'):
             with st.expander("⚠️ Last Error", expanded=True):
@@ -191,4 +212,4 @@ def render_ui(process_video_func):
                     st.rerun()
     except Exception as e:
         logger.error(f"[UI] Render UI error: {e}", exc_info=True)
-        st.error(f"UI rendering error: {str(e)}. Check logs for details.") 
+        st.error(f"UI rendering error: {str(e)}. Check logs for details.")