Update ui_components.py

ui_components.py

@@ -1,9 +1,7 @@
 #!/usr/bin/env python3
 """
 FIXED UI Components for Video Background Replacement
-Updated to use the fixed core processing functions with working SAM2 + MatAnyone
-Now includes TWO-STAGE processing option
-ENHANCED: Added debug logging and cancel functionality
+Updated with enhanced debug logging for two-stage checkbox issue
 """
 
 import gradio as gr
@@ -62,6 +60,7 @@ def add_debug_log(message):
     if len(DEBUG_LOG) > MAX_DEBUG_LINES:
         DEBUG_LOG = DEBUG_LOG[-MAX_DEBUG_LINES:]
     logger.info(message)
+    print(f"DEBUG: {log_entry}")  # Also print to console
     return "\n".join(DEBUG_LOG)
 
 def clear_debug_log():
@@ -77,24 +76,6 @@ def get_debug_log():
 # ============================================================================ #
 # UI HELPER FUNCTIONS
 # ============================================================================ #
-def generate_ai_background(prompt, style):
-    """Generate AI background from prompt"""
-    if not prompt or not prompt.strip():
-        return None, "Please enter a prompt", get_debug_log()
-    try:
-        add_debug_log(f"Generating AI background: {prompt[:50]}...")
-        bg_image = create_procedural_background(prompt, style, 1920, 1080)
-        if bg_image is not None:
-            with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
-                cv2.imwrite(tmp.name, bg_image)
-                add_debug_log(f"AI background saved to: {tmp.name}")
-                return tmp.name, f"Background generated: {prompt[:50]}...", get_debug_log()
-        return None, "Generation failed, try different prompt", get_debug_log()
-    except Exception as e:
-        error_msg = f"AI generation error: {str(e)}"
-        add_debug_log(error_msg)
-        return None, error_msg, get_debug_log()
-
 def switch_background_method(method):
     """Switch between background method UI groups"""
     add_debug_log(f"Switched to background method: {method}")
@@ -106,7 +87,7 @@ def switch_background_method(method):
     )
 
 def update_cache_status():
-    """Update cache status display - Simplified to remove misleading info"""
+    """Update cache status display"""
     try:
         cache = get_cache_status()
         two_stage_status = "✅ Available" if cache.get('two_stage_available', False) else "❌ Not available"
@@ -114,6 +95,11 @@ def update_cache_status():
     except Exception as e:
         return f"Status check error: {str(e)}"
 
+def on_two_stage_change(value):
+    """Debug handler for two-stage checkbox changes"""
+    add_debug_log(f"Two-stage checkbox changed to: {value} (type: {type(value)})")
+    return gr.update(visible=value)
+
 # ============================================================================ #
 # PROGRESS WRAPPER
 # ============================================================================ #
@@ -140,6 +126,15 @@ def process_video_enhanced_fixed(
         add_debug_log("ERROR: No video file provided")
         return None, "No video file provided.", get_debug_log()
     
+    # CRITICAL DEBUG: Log exact checkbox value received
+    add_debug_log(f"CHECKBOX RAW VALUE: use_two_stage = {use_two_stage}")
+    add_debug_log(f"CHECKBOX TYPE: {type(use_two_stage)}")
+    add_debug_log(f"CHECKBOX BOOL CONVERSION: {bool(use_two_stage)}")
+    
+    # Force boolean conversion to be safe
+    use_two_stage = bool(use_two_stage)
+    add_debug_log(f"FINAL use_two_stage value: {use_two_stage}")
+    
     def progress_callback(pct, desc):
         gradio_progress_wrapper(progress, pct, desc)
     
@@ -149,7 +144,8 @@ def progress_callback(pct, desc):
         add_debug_log("Starting new video processing...")
         add_debug_log(f"Video: {video_path}")
         add_debug_log(f"Background method: {bg_method}")
-        add_debug_log(f"Two-stage mode: {use_two_stage}")
+        add_debug_log(f"Two-stage mode requested: {use_two_stage}")
+        add_debug_log(f"Chroma preset: {chroma_preset}")
         
         mode_msg = "TWO-STAGE Green Screen" if use_two_stage else "Single-Stage"
         add_debug_log(f"Processing mode: {mode_msg}")
@@ -157,18 +153,24 @@ def progress_callback(pct, desc):
         if bg_method == "upload":
             if custom_img and os.path.exists(custom_img):
                 add_debug_log(f"Using uploaded image: {custom_img}")
-                result = process_video_fixed(video_path, "custom", custom_img, 
-                                         progress_callback, use_two_stage=use_two_stage,
-                                         chroma_preset=chroma_preset)
+                result = process_video_fixed(
+                    video_path, "custom", custom_img, 
+                    progress_callback, 
+                    use_two_stage=use_two_stage,
+                    chroma_preset=chroma_preset
+                )
                 return result[0], result[1], get_debug_log()
             return None, "No image uploaded. Please upload a background image.", get_debug_log()
             
         elif bg_method == "professional":
             if prof_choice and prof_choice in PROFESSIONAL_BACKGROUNDS:
                 add_debug_log(f"Using professional background: {prof_choice}")
-                result = process_video_fixed(video_path, prof_choice, None, 
-                                         progress_callback, use_two_stage=use_two_stage,
-                                         chroma_preset=chroma_preset)
+                result = process_video_fixed(
+                    video_path, prof_choice, None, 
+                    progress_callback, 
+                    use_two_stage=use_two_stage,
+                    chroma_preset=chroma_preset
+                )
                 return result[0], result[1], get_debug_log()
             return None, f"Invalid professional background: {prof_choice}", get_debug_log()
             
@@ -190,9 +192,12 @@ def progress_callback(pct, desc):
                 cv2.imwrite(temp_path, gradient_bg)
                 add_debug_log(f"Gradient saved to: {temp_path}")
                 
-                result = process_video_fixed(video_path, "custom", temp_path, 
-                                         progress_callback, use_two_stage=use_two_stage,
-                                         chroma_preset=chroma_preset)
+                result = process_video_fixed(
+                    video_path, "custom", temp_path, 
+                    progress_callback, 
+                    use_two_stage=use_two_stage,
+                    chroma_preset=chroma_preset
+                )
                 return result[0], result[1], get_debug_log()
             except Exception as e:
                 error_msg = f"Error creating gradient: {str(e)}"
@@ -202,9 +207,12 @@ def progress_callback(pct, desc):
         elif bg_method == "ai":
             if ai_img and os.path.exists(ai_img):
                 add_debug_log(f"Using AI generated image: {ai_img}")
-                result = process_video_fixed(video_path, "custom", ai_img, 
-                                         progress_callback, use_two_stage=use_two_stage,
-                                         chroma_preset=chroma_preset)
+                result = process_video_fixed(
+                    video_path, "custom", ai_img, 
+                    progress_callback, 
+                    use_two_stage=use_two_stage,
+                    chroma_preset=chroma_preset
+                )
                 return result[0], result[1], get_debug_log()
             return None, "No AI background generated. Click 'Generate Background' first.", get_debug_log()
             
@@ -233,7 +241,7 @@ def create_interface():
     """Create the complete Gradio interface with debug capabilities"""
     
     with gr.Blocks(
-        title="FIXED High-Quality Video Background Replacement",
+        title="Video Background Replacement",
         theme=gr.themes.Soft(),
         css="""
         .gradio-container { max-width: 1200px !important; }
@@ -246,6 +254,8 @@ def create_interface():
             font-size: 12px;
             padding: 10px;
             border-radius: 4px;
+            max-height: 400px;
+            overflow-y: auto;
         }
         .two-stage-box { 
             background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
@@ -256,15 +266,15 @@ def create_interface():
         }
         """
     ) as demo:
-        gr.Markdown("# FIXED Cinema-Quality Video Background Replacement")
-        gr.Markdown("**Upload a video → Choose a background → Get professional results with VALIDATED AI**")
+        gr.Markdown("# 🎬 Video Background Replacement")
+        gr.Markdown("Professional quality background replacement using AI segmentation")
         
-        # Two-stage mode banner if available
+        # Two-stage mode banner
         if TWO_STAGE_AVAILABLE:
             gr.HTML("""
             <div class="two-stage-box">
-                <strong>🎬 NEW: Two-Stage Green Screen Mode Available!</strong><br>
-                Professional cinema-grade chroma key processing for perfect edges
+                <strong>🎬 Two-Stage Green Screen Mode Available!</strong><br>
+                Cinema-grade processing with chroma key for perfect edges
             </div>
             """)
         
@@ -272,36 +282,36 @@ def create_interface():
 
         with gr.Row():
             with gr.Column(scale=1):
-                gr.Markdown("### Step 1: Upload Your Video")
-                gr.Markdown("*Supports MP4, MOV, AVI, and other common formats*")
+                gr.Markdown("### 📹 Step 1: Upload Your Video")
                 video_input = gr.Video(label="Drop your video here", height=300)
 
-                gr.Markdown("### Step 2: Choose Background Method")
-                gr.Markdown("*Select your preferred background creation method*")
+                gr.Markdown("### 🎨 Step 2: Choose Background")
                 background_method = gr.Radio(
                     choices=["upload", "professional", "colors", "ai"],
                     value="professional",
-                    label="Background Method"
+                    label="Background Method",
+                    info="Select how you want to create your background"
                 )
 
+                # Background method groups
                 with gr.Group(visible=False) as upload_group:
-                    gr.Markdown("**Upload Your Background Image**")
-                    custom_background = gr.Image(label="Drop your background image here", type="filepath")
+                    gr.Markdown("**Upload Custom Background**")
+                    custom_background = gr.Image(label="Upload image", type="filepath")
 
                 with gr.Group(visible=True) as professional_group:
-                    gr.Markdown("**Professional Background Presets**")
+                    gr.Markdown("**Professional Presets**")
                     professional_choice = gr.Dropdown(
                         choices=list(PROFESSIONAL_BACKGROUNDS.keys()),
                         value="office_modern",
-                        label="Select Professional Background"
+                        label="Select Background"
                     )
 
                 with gr.Group(visible=False) as colors_group:
-                    gr.Markdown("**Custom Colors & Gradients**")
+                    gr.Markdown("**Custom Gradient**")
                     gradient_type = gr.Dropdown(
-                        choices=["solid", "vertical", "horizontal", "diagonal", "radial", "soft_radial"],
+                        choices=["solid", "vertical", "horizontal", "diagonal", "radial"],
                         value="vertical",
-                        label="Gradient Type"
+                        label="Type"
                     )
                     with gr.Row():
                         color1 = gr.ColorPicker(label="Color 1", value="#3498db")
@@ -311,127 +321,118 @@ def create_interface():
                         use_third_color = gr.Checkbox(label="Use 3rd color", value=False)
 
                 with gr.Group(visible=False) as ai_group:
-                    gr.Markdown("**AI Generated Background**")
+                    gr.Markdown("**AI Generated (Experimental)**")
                     ai_prompt = gr.Textbox(
-                        label="Describe your background",
-                        placeholder="e.g., 'modern office with plants', 'sunset over mountains'",
+                        label="Describe background",
+                        placeholder="e.g., 'modern office with plants'",
                         lines=2
                     )
                     ai_style = gr.Dropdown(
-                        choices=["photorealistic", "artistic", "abstract", "minimalist", "corporate", "nature"],
+                        choices=["photorealistic", "artistic", "abstract"],
                         value="photorealistic",
                         label="Style"
                     )
                     with gr.Row():
-                        generate_ai_btn = gr.Button("Generate Background", variant="secondary")
-                        ai_generated_image = gr.Image(label="Generated Background", type="filepath", visible=False)
+                        generate_ai_btn = gr.Button("Generate", variant="secondary")
+                        ai_generated_image = gr.Image(label="Generated", type="filepath", visible=False)
 
-                # Wire up background method switching
+                # Wire up background switching
                 background_method.change(
                     fn=switch_background_method,
                     inputs=background_method,
                     outputs=[upload_group, professional_group, colors_group, ai_group]
                 )
 
-                gr.Markdown("### Processing Options")
+                gr.Markdown("### ⚙️ Processing Options")
                 
-                # Two-stage processing options
-                with gr.Accordion("🎬 Advanced: Two-Stage Processing", open=False):
+                # Two-stage processing
+                with gr.Accordion("Advanced Settings", open=True):
                     use_two_stage = gr.Checkbox(
-                        label="Enable Two-Stage Green Screen Processing",
+                        label="Enable Two-Stage Processing (Green Screen)",
                         value=False,
-                        info="Creates green screen intermediate for cinema-quality edges"
+                        info="Better quality but slower - creates green screen intermediate"
                     )
+                    
                     chroma_preset = gr.Dropdown(
                         choices=list(CHROMA_PRESETS.keys()) if TWO_STAGE_AVAILABLE else ["standard"],
                         value="standard",
-                        label="Chroma Key Preset",
-                        info="Edge quality: standard=balanced, tight=sharp, soft=smooth",
+                        label="Edge Quality",
+                        info="Standard=balanced, Tight=sharp edges, Soft=smooth blend",
                         visible=False
                     )
                     
+                    # Debug checkbox changes
                     use_two_stage.change(
-                        fn=lambda x: gr.update(visible=x),
+                        fn=on_two_stage_change,
                         inputs=use_two_stage,
                         outputs=chroma_preset
                     )
 
-                gr.Markdown("### Processing Controls")
+                gr.Markdown("### 🚀 Process Video")
                 with gr.Row():
-                    load_models_btn = gr.Button("Step 1: Load Models", variant="secondary")
-                    process_btn = gr.Button("Step 2: Process Video", variant="primary")
-                    cancel_btn = gr.Button("Cancel", variant="stop")
+                    load_models_btn = gr.Button("Load Models", variant="secondary")
+                    process_btn = gr.Button("Process Video", variant="primary", scale=2)
 
                 status_text = gr.Textbox(
-                    label="System Status", 
-                    value="Ready - click 'Load Models' to start",
-                    interactive=False, 
-                    lines=3,
-                    elem_classes=["status-box"]
-                )
-
-                validation_status = gr.Textbox(
-                    label="Processing Mode",
-                    value=update_cache_status(),
+                    label="Status", 
+                    value="Ready - Click 'Load Models' first",
                     interactive=False,
-                    lines=1,
-                    elem_classes=["status-box"]
+                    lines=2
                 )
 
             with gr.Column(scale=1):
-                gr.Markdown("### Your Results")
-                video_output = gr.Video(label="Your Processed Video", height=400)
+                gr.Markdown("### 🎥 Result")
+                video_output = gr.Video(label="Processed Video", height=400)
+                
                 result_text = gr.Textbox(
-                    label="Processing Results",
+                    label="Processing Info",
                     interactive=False,
-                    lines=8,
-                    placeholder="Processing status will appear here...",
-                    elem_classes=["status-box"]
+                    lines=6,
+                    placeholder="Results will appear here..."
                 )
                 
-                # Debug panel
-                with gr.Accordion("🔧 Debug Log", open=False):
+                # Debug panel - OPEN by default for testing
+                with gr.Accordion("🔧 Debug Log", open=True):
                     debug_output = gr.Textbox(
                         label="Debug Output",
                         value=get_debug_log(),
-                        lines=15,
+                        lines=12,
                         max_lines=20,
                         interactive=False,
                         elem_classes=["debug-box"]
                     )
                     with gr.Row():
-                        refresh_debug_btn = gr.Button("Refresh Log", size="sm")
-                        clear_debug_btn = gr.Button("Clear Log", size="sm")
+                        refresh_debug_btn = gr.Button("Refresh", size="sm")
+                        clear_debug_btn = gr.Button("Clear", size="sm")
 
         # Event handlers
-        load_event = load_models_btn.click(
+        load_models_btn.click(
             fn=load_models_with_progress_fixed, 
             outputs=[status_text, debug_output]
         )
         
-        generate_ai_btn.click(
-            fn=generate_ai_background, 
-            inputs=[ai_prompt, ai_style], 
-            outputs=[ai_generated_image, status_text, debug_output]
-        )
-        
-        # Process with cancel capability
-        process_event = process_btn.click(
+        # Main processing - ensure all inputs are connected
+        process_btn.click(
             fn=process_video_enhanced_fixed,
-            inputs=[video_input, background_method, custom_background, professional_choice,
-                    gradient_type, color1, color2, color3, use_third_color,
-                    ai_prompt, ai_style, ai_generated_image,
-                    use_two_stage, chroma_preset],
+            inputs=[
+                video_input,           # video_path
+                background_method,     # bg_method
+                custom_background,     # custom_img
+                professional_choice,   # prof_choice
+                gradient_type,        # grad_type
+                color1,               # color1
+                color2,               # color2
+                color3,               # color3
+                use_third_color,      # use_third
+                ai_prompt,            # ai_prompt
+                ai_style,             # ai_style
+                ai_generated_image,   # ai_img
+                use_two_stage,        # use_two_stage - THIS IS THE KEY INPUT
+                chroma_preset         # chroma_preset
+            ],
             outputs=[video_output, result_text, debug_output]
         )
         
-        # Cancel button
-        cancel_btn.click(
-            fn=lambda: (add_debug_log("Processing cancelled by user"), "Processing cancelled")[1],
-            outputs=[status_text],
-            cancels=[process_event, load_event]
-        )
-        
         # Debug controls
         refresh_debug_btn.click(
             fn=get_debug_log,
@@ -443,20 +444,17 @@ def create_interface():
             outputs=[debug_output]
         )
 
-        # Auto-update validation status
-        load_models_btn.click(
-            fn=update_cache_status,
-            outputs=[validation_status]
-        )
-
-        with gr.Accordion("Quality & Features", open=False):
+        # Info section
+        with gr.Accordion("ℹ️ Information", open=False):
             gr.Markdown("""
-            ### Features:
-            - Single-Stage: Fast processing with good quality
-            - Two-Stage: Cinema-grade green screen processing
-            - Debug logging for troubleshooting
-            - Cancel button to stop long processes
-            - SAM2 + MatAnyone AI models
+            ### Processing Modes:
+            - **Single-Stage**: Fast, good quality for most videos
+            - **Two-Stage**: Slower but cinema-quality edges using green screen
+            
+            ### Tips:
+            - Start with 720p videos for testing
+            - Two-stage mode works best with clear subjects
+            - Processing takes 2-4 minutes for 1-minute videos
             """)
 
     return demo