Update ui_components.py

ui_components.py

@@ -2,6 +2,7 @@
 """
 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
 """
 
 import gradio as gr
@@ -13,9 +14,17 @@
 from typing import Optional
 
 # Import FIXED core processing functions
-from app import load_models_with_validation, process_video_fixed, get_model_status
+from app import load_models_with_validation, process_video_fixed, get_model_status, get_cache_status
 from utilities import PROFESSIONAL_BACKGROUNDS, create_procedural_background
 
+# Check if two-stage is available
+try:
+    from two_stage_processor import CHROMA_PRESETS
+    TWO_STAGE_AVAILABLE = True
+except ImportError:
+    TWO_STAGE_AVAILABLE = False
+    CHROMA_PRESETS = {'standard': {}}
+
 logger = logging.getLogger(__name__)
 
 # ============================================================================ #
@@ -66,12 +75,17 @@ def update_cache_status():
     """Update cache status display with actual model state - FIXED"""
     try:
         status = get_model_status()
+        cache = get_cache_status()
         
         sam2_status = f"✅ Loaded & Validated" if status['sam2'] == 'Ready' else "❌ Not loaded"
         matanyone_status = f"✅ Loaded & Validated" if status['matanyone'] == 'Ready' else "❌ Not loaded"
         models_status = "✅ Models Validated" if status['validated'] else "❌ Not validated"
+        two_stage_status = "✅ Available" if cache.get('two_stage_available', False) else "❌ Not available"
         
-        return f"SAM2: {sam2_status}\nMatAnyone: {matanyone_status}\nStatus: {models_status}"
+        return (f"SAM2: {sam2_status}\n"
+                f"MatAnyone: {matanyone_status}\n"
+                f"Two-Stage: {two_stage_status}\n"
+                f"Status: {models_status}")
     except Exception as e:
         return f"Status check error: {str(e)}"
 
@@ -87,14 +101,15 @@ def gradio_progress_wrapper(progress_obj, pct, desc):
         pass
 
 # ============================================================================ #
-# MAIN PROCESSING FUNCTION FOR UI (FIXED)
+# MAIN PROCESSING FUNCTION FOR UI (FIXED WITH TWO-STAGE)
 # ============================================================================ #
 def process_video_enhanced_fixed(
     video_path, bg_method, custom_img, prof_choice, grad_type,
     color1, color2, color3, use_third, ai_prompt, ai_style, ai_img,
+    use_two_stage, chroma_preset,  # NEW PARAMETERS
     progress: Optional[gr.Progress] = None
 ):
-    """FIXED enhanced video processing function that uses validated SAM2 + MatAnyone"""
+    """FIXED enhanced video processing function with two-stage option"""
     if not video_path:
         return None, "No video file provided."
     
@@ -102,14 +117,22 @@ def progress_callback(pct, desc):
         gradio_progress_wrapper(progress, pct, desc)
     
     try:
+        # Determine the processing mode message
+        mode_msg = "TWO-STAGE Green Screen" if use_two_stage else "Single-Stage"
+        logger.info(f"Processing with {mode_msg} mode")
+        
         if bg_method == "upload":
             if custom_img and os.path.exists(custom_img):
-                return process_video_fixed(video_path, "custom", custom_img, progress_callback)
+                return process_video_fixed(video_path, "custom", custom_img, 
+                                         progress_callback, use_two_stage=use_two_stage,
+                                         chroma_preset=chroma_preset)
             return None, "No image uploaded. Please upload a background image."
             
         elif bg_method == "professional":
             if prof_choice and prof_choice in PROFESSIONAL_BACKGROUNDS:
-                return process_video_fixed(video_path, prof_choice, None, progress_callback)
+                return process_video_fixed(video_path, prof_choice, None, 
+                                         progress_callback, use_two_stage=use_two_stage,
+                                         chroma_preset=chroma_preset)
             return None, f"Invalid professional background: {prof_choice}"
             
         elif bg_method == "colors":
@@ -127,13 +150,17 @@ def progress_callback(pct, desc):
                 gradient_bg = create_professional_background(bg_config, 1920, 1080)
                 temp_path = f"/tmp/gradient_{int(time.time())}.png"
                 cv2.imwrite(temp_path, gradient_bg)
-                return process_video_fixed(video_path, "custom", temp_path, progress_callback)
+                return process_video_fixed(video_path, "custom", temp_path, 
+                                         progress_callback, use_two_stage=use_two_stage,
+                                         chroma_preset=chroma_preset)
             except Exception as e:
                 return None, f"Error creating gradient: {str(e)}"
                 
         elif bg_method == "ai":
             if ai_img and os.path.exists(ai_img):
-                return process_video_fixed(video_path, "custom", ai_img, progress_callback)
+                return process_video_fixed(video_path, "custom", ai_img, 
+                                         progress_callback, use_two_stage=use_two_stage,
+                                         chroma_preset=chroma_preset)
             return None, "No AI background generated. Click 'Generate Background' first."
             
         else:
@@ -151,10 +178,10 @@ def progress_callback(pct, desc):
     return load_models_with_validation(progress_callback)
 
 # ============================================================================ #
-# MAIN INTERFACE CREATION (FIXED)
+# MAIN INTERFACE CREATION (FIXED WITH TWO-STAGE)
 # ============================================================================ #
 def create_interface():
-    """Create the complete Gradio interface with FIXED processing"""
+    """Create the complete Gradio interface with FIXED processing and TWO-STAGE option"""
     
     with gr.Blocks(
         title="FIXED High-Quality Video Background Replacement",
@@ -163,11 +190,28 @@ def create_interface():
         .gradio-container { max-width: 1200px !important; }
         .progress-bar { background: linear-gradient(90deg, #3498db, #2ecc71) !important; }
         .status-box { background-color: #f8f9fa; border: 1px solid #dee2e6; border-radius: 6px; }
+        .two-stage-box { 
+            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+            color: white;
+            padding: 10px;
+            border-radius: 8px;
+            margin: 10px 0;
+        }
         """
     ) 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("*Powered by PROPERLY INTEGRATED SAM2 + MatAnyone with comprehensive validation*")
+        
+        # Two-stage mode banner if available
+        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
+            </div>
+            """)
+        
         gr.Markdown("---")
 
         with gr.Row():
@@ -239,11 +283,35 @@ def create_interface():
                     outputs=[upload_group, professional_group, colors_group, ai_group]
                 )
 
+                gr.Markdown("### Processing Options")
+                
+                # Two-stage processing options
+                with gr.Accordion("🎬 Advanced: Two-Stage Processing", open=False):
+                    use_two_stage = gr.Checkbox(
+                        label="Enable Two-Stage Green Screen Processing",
+                        value=False,
+                        info="Creates green screen intermediate for cinema-quality edges"
+                    )
+                    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",
+                        visible=False
+                    )
+                    
+                    # Show/hide chroma preset based on two-stage checkbox
+                    use_two_stage.change(
+                        fn=lambda x: gr.update(visible=x),
+                        inputs=use_two_stage,
+                        outputs=chroma_preset
+                    )
+
                 gr.Markdown("### Processing Controls")
                 gr.Markdown("*First load and validate the AI models, then process your video*")
                 with gr.Row():
                     load_models_btn = gr.Button("Step 1: Load & Validate AI Models", variant="secondary")
-                    process_btn = gr.Button("Step 2: Process Video (FIXED)", variant="primary")
+                    process_btn = gr.Button("Step 2: Process Video", variant="primary")
 
                 status_text = gr.Textbox(
                     label="System Status", 
@@ -258,19 +326,19 @@ def create_interface():
                     label="AI Models Validation Status",
                     value=update_cache_status(),
                     interactive=False,
-                    lines=3,
+                    lines=4,
                     elem_classes=["status-box"]
                 )
 
             with gr.Column(scale=1):
                 gr.Markdown("### Your Results")
-                gr.Markdown("*FIXED processed video with validated SAM2 + MatAnyone will appear here*")
-                video_output = gr.Video(label="Your FIXED Processed Video", height=400)
+                gr.Markdown("*Processed video with validated SAM2 + MatAnyone will appear here*")
+                video_output = gr.Video(label="Your Processed Video", height=400)
                 result_text = gr.Textbox(
                     label="Processing Results",
                     interactive=False,
-                    lines=8,
-                    placeholder="FIXED processing status and detailed results will appear here...",
+                    lines=10,
+                    placeholder="Processing status and detailed results will appear here...",
                     elem_classes=["status-box"]
                 )
 
@@ -305,11 +373,13 @@ def create_interface():
             outputs=[ai_generated_image, status_text]
         )
         
+        # Updated process button with two-stage parameters
         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],
+                    ai_prompt, ai_style, ai_generated_image,
+                    use_two_stage, chroma_preset],  # Added two-stage parameters
             outputs=[video_output, result_text]
         )
 
@@ -319,19 +389,31 @@ def create_interface():
             outputs=[validation_status]
         )
 
-        with gr.Accordion("FIXED Quality & Features", open=False):
+        with gr.Accordion("Quality & Features", open=False):
             gr.Markdown("""
-            ### FIXED Single-Stage Cinema-Quality Features:
-            **Processing**: Original → Final Background (VALIDATED SAM2 + MatAnyone)
-            **Quality**: Comprehensive validation, enhanced fallbacks, edge feathering, gamma correction
-            **Models**: SAM2 with functionality testing + MatAnyone with proper error handling
-            **Segmentation**: Multi-point SAM2 strategy with advanced OpenCV fallback
-            **Refinement**: MatAnyone refinement with enhanced OpenCV backup
-            **Architecture**: Separated UI components with proper error handling and validation
-            **Output**: H.264 CRF 18, AAC 192kbps with preserved audio
+            ### Cinema-Quality Features:
+            
+            **Single-Stage Mode (Default):**
+            - Original → Final Background
+            - Fast processing with good quality
+            - Edge threshold: 140, feather: 1px
+            - Morphological edge cleaning
+            
+            **Two-Stage Mode (Professional):**
+            - Stage 1: Original → Green Screen
+            - Stage 2: Green Screen → Final (Chroma Key)
+            - Perfect edges with spill suppression
+            - Adjustable tolerance and softness
+            - Cinema-grade quality
+            
+            **Core Technology:**
+            - SAM2 with multi-point strategy
+            - MatAnyone refinement every 3 frames
+            - Enhanced OpenCV fallbacks
+            - H.264 CRF 18, AAC 192kbps audio
             """)
 
         gr.Markdown("---")
-        gr.Markdown("*FIXED Cinema-Quality Video Background Replacement — Validated SAM2 + MatAnyone pipeline*")
+        gr.Markdown("*Cinema-Quality Video Background Replacement — Validated SAM2 + MatAnyone pipeline*")
 
     return demo