#!/usr/bin/env python3 """ UI Components for BackgroundFX Pro - with Start/Stop, frame counter, and live FPS """ import gradio as gr import time import threading # --- Import backend logic --- try: from core.app import ( processor, load_models_with_validation, process_video_fixed, get_model_status, get_cache_status, ) CORE_FUNCTIONS_AVAILABLE = True except Exception as e: print(f"[UI] Core functions import failed: {e}") CORE_FUNCTIONS_AVAILABLE = False try: from utils.backgrounds import PROFESSIONAL_BACKGROUNDS UTILITIES_AVAILABLE = True except Exception as e: PROFESSIONAL_BACKGROUNDS = {"office_modern": {"name": "Modern Office", "description": "Default office background"}} UTILITIES_AVAILABLE = False # --- UI state management --- stop_event = threading.Event() def create_interface(): is_processing = gr.State(False) frame_progress = gr.State({"current": 0, "total": 0, "fps": 0.0}) def enhanced_process_video( video_path, bg_method, custom_img, prof_choice, use_two_stage, chroma_preset, quality_preset, is_processing_state, frame_progress_state, progress=gr.Progress() ): """ Handles video processing (start/stop) and live frame/fps reporting. """ # If already processing: user wants to STOP if is_processing_state: stop_event.set() return None, "Processing stopped by user.", "Processing stopped.", False, {"current": 0, "total": 0, "fps": 0.0} # Otherwise: START processing stop_event.clear() is_processing_state = True frame_progress_state = {"current": 0, "total": 0, "fps": 0.0} last_update = time.time() last_frame = 0 if not CORE_FUNCTIONS_AVAILABLE: return None, "Core backend not available.", "System error.", False, frame_progress_state if not processor.models_loaded: return None, "Models not loaded.", "Load models first.", False, frame_progress_state if not video_path: return None, "No video uploaded.", "Upload a video.", False, frame_progress_state # Choose background if bg_method == "professional" and not prof_choice: return None, "No professional background selected.", "Choose a background.", False, frame_progress_state if bg_method == "upload" and not custom_img: return None, "No custom background image.", "Upload a background image.", False, frame_progress_state try: def progress_callback(pct, desc, current_frame=None, total_frames=None): # User cancellation check if stop_event.is_set(): raise Exception("Processing stopped by user.") # Standard Gradio progress if progress: progress(pct, desc) # Update frame count/fps nonlocal last_update, last_frame, frame_progress_state now = time.time() fps = 0.0 if current_frame is not None and total_frames is not None: dt = now - last_update if last_update else 0.01 frames_done = current_frame - last_frame if last_frame is not None else 0 fps = frames_done / dt if dt > 0 else 0.0 last_update = now last_frame = current_frame frame_progress_state = { "current": int(current_frame), "total": int(total_frames), "fps": round(fps, 2) } return desc result_path, result_message = process_video_fixed( video_path=video_path, background_choice=prof_choice if bg_method == "professional" else "custom", custom_background_path=custom_img if bg_method == "upload" else None, progress_callback=lambda pct, desc: progress_callback(pct, desc), # You may want to pass current_frame/total_frames use_two_stage=bool(use_two_stage), chroma_preset=chroma_preset or "standard", preview_mask=False, preview_greenscreen=False, stop_event=stop_event, # <-- Backend must support this argument! ) is_processing_state = False return result_path, result_message, "Processing completed.", False, frame_progress_state except Exception as e: is_processing_state = False return None, f"Error: {str(e)}", "Error during processing.", False, frame_progress_state def handle_model_loading(progress=gr.Progress()): """ Wrapper function for model loading with proper progress handling """ try: # Call the model loading function with progress callback result = load_models_with_validation(progress) return result except Exception as e: return f"Model loading failed: {str(e)}" # --- UI setup --- with gr.Blocks( title="BackgroundFX Pro", css=""" .main-header { text-align: center; margin-bottom: 20px; } .status-box { background: #f8f9fa; padding: 15px; border-radius: 8px; margin: 10px 0; } .live-progress { font-size: 1.1em; font-weight: bold; color: #004085; } """ ) as demo: with gr.Row(): gr.Markdown("# BackgroundFX Pro - Video Background Replacement", elem_classes=["main-header"]) # --- Main UI fields --- with gr.Row(): with gr.Column(): video_input = gr.Video(label="Upload Video", height=280) background_method = gr.Radio( ["professional", "upload"], value="professional", label="Background Method" ) professional_choice = gr.Dropdown( choices=[(bg['name'], key) for key, bg in PROFESSIONAL_BACKGROUNDS.items()], value=list(PROFESSIONAL_BACKGROUNDS.keys())[0], label="Professional Background" ) custom_background = gr.Image(label="Upload Background", type="filepath", visible=False) quality_preset = gr.Dropdown( choices=[ ("Fast", "fast"), ("Balanced", "balanced"), ("High Quality", "high") ], value="balanced", label="Quality Preset" ) use_two_stage = gr.Checkbox("Enable Two-Stage Mode", value=False) chroma_preset = gr.Dropdown( choices=[("Standard", "standard"), ("Studio", "studio"), ("Outdoor", "outdoor")], value="standard", label="Chroma Key Preset" ) # Dynamic show/hide def update_background_visibility(method): return ( gr.update(visible=(method == "professional")), gr.update(visible=(method == "upload")) ) background_method.change( fn=update_background_visibility, inputs=background_method, outputs=[professional_choice, custom_background] ) with gr.Row(): load_models_btn = gr.Button("Load Models", variant="secondary") process_btn = gr.Button("Process Video", variant="primary") status_text = gr.Textbox(label="Status", value="Ready", interactive=False) frame_info = gr.Textbox(label="Frame Progress", value="", interactive=False) fps_info = gr.Textbox(label="Frames/sec", value="", interactive=False) with gr.Column(): video_output = gr.Video(label="Processed Video", height=360) result_info = gr.Textbox(label="Processing Info", lines=8, interactive=False) debug_info = gr.Textbox(label="Debug Log", lines=6, interactive=False, visible=False) # --- Button actions --- def update_process_btn(is_processing): return gr.update(value="Stop Processing" if is_processing else "Process Video") is_processing.change(update_process_btn, inputs=is_processing, outputs=process_btn) def update_frame_fields(frame_progress): total = frame_progress.get("total", 0) current = frame_progress.get("current", 0) fps = frame_progress.get("fps", 0.0) txt = f"{current}/{total}" if total else "" return txt, f"{fps:.2f} FPS" if fps else "" frame_progress.change(update_frame_fields, inputs=frame_progress, outputs=[frame_info, fps_info]) process_btn.click( fn=enhanced_process_video, inputs=[ video_input, background_method, custom_background, professional_choice, use_two_stage, chroma_preset, quality_preset, is_processing, frame_progress ], outputs=[video_output, result_info, debug_info, is_processing, frame_progress], show_progress=True ) # Fixed model loading button - let Gradio handle progress automatically load_models_btn.click( fn=handle_model_loading, outputs=[status_text], show_progress=True ) return demo def create_ui(): return create_interface()