Update app.py

app.py

@@ -1,295 +1,305 @@
-#!/usr/bin/env python3
-"""
-MyAvatar Video Background Replacer - Streamlit UI
-Main interface for two-stage video processing pipeline
-"""
-
 import streamlit as st
-import sys
+import cv2
+import numpy as np
+import torch
+import tempfile
+import os
+import logging
 from pathlib import Path
-from PIL import Image
-
-# Add project root to path
-sys.path.append(str(Path(__file__).parent.absolute()))
-
-# Import processing modules
-from model_loaders import load_sam2_predictor, load_matanyone_processor, clear_model_cache, get_memory_usage
-from video_pipeline import stage1_create_transparent_video, stage2_composite_background
+import time
 
-# Persistent temp dir
-TMP_DIR = Path("tmp")
-TMP_DIR.mkdir(parents=True, exist_ok=True)
-
-# Page config
-st.set_page_config(
-    page_title="MyAvatar - Video Background Replacer",
-    page_icon="🎥",
-    layout="wide",
-    initial_sidebar_state="expanded"
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='[%(asctime)s] %(levelname)s: %(message)s',
+    datefmt='%H:%M:%S'
 )
+logger = logging.getLogger(__name__)
 
-# Styling
-def add_logo():
-    st.markdown(
-        """
-        <style>
-            .main .block-container { padding-top: 2rem; padding-bottom: 2rem; }
-            .stButton>button { width: 100%; background-color: #4CAF50; color: white; font-weight: bold; transition: all 0.3s; }
-            .stButton>button:hover { background-color: #45a049; }
-            .stProgress > div > div > div > div { background-color: #4CAF50; }
-            .stage-indicator { background: linear-gradient(90deg, #4CAF50, #45a049); color: white; padding: 10px; border-radius: 5px; margin: 10px 0; text-align: center; font-weight: bold; }
-        </style>
-        """,
-        unsafe_allow_html=True
-    )
+# Import model loaders
+from model_loaders import load_sam2, load_matanyone, pose
 
-def show_memory_info():
-    """Display memory usage in sidebar with model testing."""
-    memory_info = get_memory_usage()
-    with st.sidebar:
-        st.markdown("### 🧠 Memory Usage")
-        if 'gpu_allocated' in memory_info:
-            st.metric("GPU Memory", f"{memory_info['gpu_allocated']:.1f}GB", 
-                     f"Free: {memory_info['gpu_free']:.1f}GB")
-        st.metric("RAM Usage", f"{memory_info['ram_used']:.1f}GB",
-                 f"Available: {memory_info['ram_available']:.1f}GB")
-        
-        if st.button("🧪 Test Models", help="Test if SAM2 and MatAnyone can load"):
-            with st.spinner("Testing model loading..."):
-                try:
-                    sam2_test = load_sam2_predictor()
-                    st.success("✅ SAM2 loads successfully") if sam2_test else st.error("❌ SAM2 failed to load")
-                    
-                    matanyone_test = load_matanyone_processor()
-                    st.success("✅ MatAnyone loads successfully") if matanyone_test else st.error("❌ MatAnyone failed to load")
-                except Exception as e:
-                    st.error(f"Model test failed: {e}")
-        
-        if st.button("🧹 Clear Cache", help="Free up memory by clearing model cache"):
-            clear_model_cache()
-            st.success("Cache cleared!")
-            st.experimental_rerun()
+st.set_page_config(page_title="Video Matting", layout="wide")
 
-def initialize_session_state():
-    """Initialize all session state variables."""
-    defaults = {
-        'uploaded_video': None,
-        'bg_image': None,
-        'bg_image_info': None,
-        'bg_color': "#00FF00",
-        'bg_type': "image",
-        'transparent_video_path': None,
-        'final_video_path': None,
-        'processing_stage1': False,
-        'processing_stage2': False
-    }
-    for key, value in defaults.items():
-        if key not in st.session_state:
-            st.session_state[key] = value
+def log_and_progress(progress_callback, stage, progress, message):
+    """Unified logging and progress reporting"""
+    timestamp = time.strftime("%H:%M:%S")
+    log_msg = f"[{timestamp}] Stage {stage} ({progress:.0%}): {message}"
+    logger.info(log_msg)
+    print(log_msg, flush=True)
+    progress_callback(stage, progress, message)
 
-def handle_video_upload():
-    """Handle video file upload."""
-    uploaded = st.file_uploader(
-        "📹 Upload Video",
-        type=["mp4", "mov", "avi", "mkv"],
-        key="video_uploader",
-        help="Recommended: Videos under 30 seconds for faster processing"
-    )
-    if uploaded is not None:
-        file_size_mb = uploaded.size / (1024 * 1024)
-        if file_size_mb > 100:
-            st.warning(f"⚠️ Large file detected ({file_size_mb:.1f}MB). Processing may take longer.")
-        st.session_state.uploaded_video = uploaded
-        st.session_state.transparent_video_path = None
-        st.session_state.final_video_path = None
-
-def show_video_preview():
-    """Display uploaded video preview."""
-    st.markdown("### Video Preview")
-    if st.session_state.uploaded_video is not None:
-        video_bytes = st.session_state.uploaded_video.getvalue()
-        st.video(video_bytes)
-        st.session_state.uploaded_video.seek(0)
-
-def handle_background_selection():
-    """Handle background type selection."""
-    st.markdown("### Background Options")
-    bg_type = st.radio("Select Background Type:", ["Image", "Color"], horizontal=True, key="bg_type_radio")
-    st.session_state.bg_type = bg_type.lower()
+def process_video(input_path, output_path, progress_callback):
+    """Main video processing pipeline - Stage 1: Segmentation, Stage 2: Matting"""
     
-    if bg_type == "Image":
-        handle_image_background()
-    elif bg_type == "Color":
-        handle_color_background()
-
-def handle_image_background():
-    """Handle image background upload and preview."""
-    bg_image = st.file_uploader("🖼️ Upload Background Image", type=["jpg", "png", "jpeg"], 
-                                key="bg_image_uploader", help="Recommended: Images under 5MB")
+    # ============================================================
+    # STAGE 1: SEGMENTATION (Load models, read video, segment all frames)
+    # ============================================================
+    logger.info("="*60)
+    logger.info("STAGE 1: PERSON SEGMENTATION")
+    logger.info("="*60)
+    stage1_start = time.time()
     
-    if bg_image is not None:
-        image_size_mb = bg_image.size / (1024 * 1024)
-        if image_size_mb > 10:
-            st.warning(f"⚠️ Large image ({image_size_mb:.1f}MB). Consider resizing.")
+    # 1.1: Load SAM2
+    log_and_progress(progress_callback, 1, 0.0, "Loading SAM2 model...")
+    try:
+        sam_predictor = load_sam2()
+        logger.info(f"✅ SAM2 loaded successfully")
+    except Exception as e:
+        logger.error(f"❌ SAM2 loading failed: {e}")
+        raise
+    
+    # 1.2: Load video
+    log_and_progress(progress_callback, 1, 0.1, "Opening video file...")
+    try:
+        cap = cv2.VideoCapture(input_path)
+        if not cap.isOpened():
+            raise ValueError("Failed to open video file")
         
-        current_file_info = f"{bg_image.name}_{bg_image.size}"
-        if st.session_state.bg_image_info != current_file_info:
-            st.session_state.bg_image = Image.open(bg_image)
-            st.session_state.bg_image_info = current_file_info
-            st.session_state.final_video_path = None
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         
-        if st.session_state.bg_image is not None:
-            st.image(st.session_state.bg_image, caption="Selected Background", use_container_width=True)
-    else:
-        st.session_state.bg_image = None
-        st.session_state.bg_image_info = None
-
-def handle_color_background():
-    """Handle solid color background selection."""
-    st.markdown("#### Select a Color")
-    old_color = st.session_state.get('bg_color', "#00FF00")
+        logger.info(f"📹 Video: {width}x{height} @ {fps:.2f}fps, {total_frames} frames")
+    except Exception as e:
+        logger.error(f"❌ Video opening failed: {e}")
+        raise
     
-    color_presets = {
-        "Pure White": "#FFFFFF",
-        "Pure Black": "#000000",
-        "Light Gray": "#F5F5F5",
-        "Professional Blue": "#0078D4",
-        "Corporate Green": "#107C10",
-        "Custom": old_color
-    }
+    # 1.3: Read all frames
+    log_and_progress(progress_callback, 1, 0.15, f"Reading {total_frames} frames...")
+    frames = []
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frames.append(frame)
+    cap.release()
+    logger.info(f"✅ Read {len(frames)} frames")
     
-    cols = st.columns(3)
-    for i, (name, color) in enumerate(color_presets.items()):
-        with cols[i % 3]:
-            if name == "Custom":
-                new_color = st.color_picker("Custom Color", old_color, key="custom_color_picker")
-                if new_color != old_color:
-                    st.session_state.bg_color = new_color
-                    st.session_state.final_video_path = None
+    # 1.4: Segment all frames
+    log_and_progress(progress_callback, 1, 0.2, "Starting person segmentation...")
+    masks = []
+    
+    for i, frame in enumerate(frames):
+        # Progress from 0.2 to 1.0 during segmentation
+        progress = 0.2 + (0.8 * i / len(frames))
+        
+        if i % 10 == 0 or i == 0 or i == len(frames) - 1:
+            log_and_progress(progress_callback, 1, progress, 
+                           f"Segmenting frame {i+1}/{len(frames)}...")
+        
+        try:
+            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            sam_predictor.set_image(rgb_frame)
+            
+            # Detect person using MediaPipe
+            mp_results = pose.process(rgb_frame)
+            
+            if mp_results.pose_landmarks:
+                # Get bounding box from landmarks
+                h, w = frame.shape[:2]
+                landmarks = mp_results.pose_landmarks.landmark
+                xs = [lm.x * w for lm in landmarks]
+                ys = [lm.y * h for lm in landmarks]
+                x1, x2 = max(0, min(xs) - 20), min(w, max(xs) + 20)
+                y1, y2 = max(0, min(ys) - 20), min(h, max(ys) + 20)
+                
+                input_box = np.array([x1, y1, x2, y2])
+                sam_masks, _, _ = sam_predictor.predict(
+                    point_coords=None,
+                    point_labels=None,
+                    box=input_box[None, :],
+                    multimask_output=False
+                )
+                masks.append(sam_masks[0])
             else:
-                if st.button(name, key=f"color_{name}", use_container_width=True):
-                    st.session_state.bg_color = color
-                    st.session_state.final_video_path = None
-                st.markdown(f'<div style="background-color:{color}; height:30px; border-radius:4px; margin-top:-10px;"></div>', 
-                           unsafe_allow_html=True)
-
-def main():
-    """Main application entry point."""
-    add_logo()
+                masks.append(np.zeros((h, w), dtype=bool))
+                
+        except Exception as e:
+            logger.error(f"❌ Frame {i+1} segmentation failed: {e}")
+            h, w = frame.shape[:2]
+            masks.append(np.zeros((h, w), dtype=bool))
+    
+    stage1_time = time.time() - stage1_start
+    logger.info("="*60)
+    logger.info(f"✅ STAGE 1 COMPLETE in {stage1_time:.1f}s")
+    logger.info(f"   Segmented {len(masks)} frames")
+    logger.info("="*60)
     
-    st.markdown("""
-        <div style="text-align: center; margin-bottom: 30px;">
-            <h1>🎥 Video Background Replacer</h1>
-            <p>Two-Stage Processing: SAM2 + MatAnyone → Transparent → Composite</p>
-        </div>
-    """, unsafe_allow_html=True)
-    st.markdown("---")
+    # ============================================================
+    # STAGE 2: MATTING (Refine all masks, smooth, write video)
+    # ============================================================
+    logger.info("="*60)
+    logger.info("STAGE 2: HIGH-QUALITY MATTING")
+    logger.info("="*60)
+    stage2_start = time.time()
     
-    initialize_session_state()
-    show_memory_info()
+    # 2.1: Load MatAnyone
+    log_and_progress(progress_callback, 2, 0.0, "Loading MatAnyone model...")
+    try:
+        matanyone = load_matanyone()
+        logger.info(f"✅ MatAnyone loaded successfully")
+    except Exception as e:
+        logger.error(f"❌ MatAnyone loading failed: {e}")
+        raise
     
-    col1, col2 = st.columns([1, 1], gap="large")
+    # 2.2: Process all frames with MatAnyone
+    log_and_progress(progress_callback, 2, 0.1, "Refining alpha mattes...")
+    alphas = []
     
-    # LEFT COLUMN: Video Upload & Stage 1
-    with col1:
-        st.header("1. Upload Video")
-        handle_video_upload()
-        show_video_preview()
+    for i, (frame, mask) in enumerate(zip(frames, masks)):
+        # Progress from 0.1 to 0.6 during matting
+        progress = 0.1 + (0.5 * i / len(frames))
+        
+        if i % 10 == 0 or i == 0 or i == len(frames) - 1:
+            log_and_progress(progress_callback, 2, progress,
+                           f"Matting frame {i+1}/{len(frames)}...")
         
-        st.markdown('<div class="stage-indicator">STAGE 1: Create Transparent Video</div>', unsafe_allow_html=True)
+        try:
+            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            
+            if mask.any():
+                result = matanyone(rgb_frame, mask)
+                alpha = result['alpha']
+            else:
+                alpha = np.zeros((frame.shape[0], frame.shape[1]), dtype=np.float32)
+            
+            alphas.append(alpha)
+            
+        except Exception as e:
+            logger.error(f"❌ Frame {i+1} matting failed: {e}")
+            alphas.append(np.zeros((frame.shape[0], frame.shape[1]), dtype=np.float32))
+    
+    logger.info(f"✅ Matted {len(alphas)} frames")
+    
+    # 2.3: Temporal smoothing
+    log_and_progress(progress_callback, 2, 0.65, "Applying temporal smoothing to eliminate jitter...")
+    try:
+        smoothed_alphas = []
+        window_size = 5  # 5-frame window (current + 2 before + 2 after)
+        half_window = window_size // 2
         
-        stage1_disabled = not st.session_state.uploaded_video or st.session_state.processing_stage1
+        for i in range(len(alphas)):
+            start_idx = max(0, i - half_window)
+            end_idx = min(len(alphas), i + half_window + 1)
+            window_alphas = alphas[start_idx:end_idx]
+            
+            # Average the alphas in the window
+            smoothed = np.mean(window_alphas, axis=0)
+            smoothed_alphas.append(smoothed)
         
-        if st.button("🎭 Create Transparent Video", type="primary", disabled=stage1_disabled, 
-                    use_container_width=True, help="Remove background using SAM2 + MatAnyone AI"):
-            with st.spinner("Stage 1: Creating transparent video..."):
-                st.session_state.processing_stage1 = True
-                try:
-                    transparent_path = stage1_create_transparent_video(st.session_state.uploaded_video)
-                    if transparent_path:
-                        st.session_state.transparent_video_path = transparent_path
-                        st.success("✅ Stage 1 Complete: Transparent video created!")
-                        st.balloons()
-                    else:
-                        st.error("❌ Stage 1 Failed: Could not create transparent video")
-                except Exception as e:
-                    st.error(f"❌ Stage 1 Error: {str(e)}")
-                finally:
-                    st.session_state.processing_stage1 = False
+        logger.info(f"✅ Applied {window_size}-frame temporal smoothing")
+        alphas = smoothed_alphas
         
-        # Show transparent video result
-        if st.session_state.get('transparent_video_path'):
-            st.markdown("#### Transparent Video Result")
-            try:
-                with open(st.session_state.transparent_video_path, 'rb') as f:
-                    transparent_bytes = f.read()
-                st.video(transparent_bytes)
-                st.download_button("💾 Download Transparent Video (.mov)", data=transparent_bytes,
-                                  file_name="transparent_video.mov", mime="video/quicktime",
-                                  use_container_width=True)
-                st.caption(f"Size: {len(transparent_bytes) / (1024**2):.1f}MB")
-            except Exception as e:
-                st.error(f"Error displaying transparent video: {str(e)}")
+    except Exception as e:
+        logger.error(f"⚠️ Smoothing failed: {e}, using unsmoothed alphas")
     
-    # RIGHT COLUMN: Background Selection & Stage 2
-    with col2:
-        st.header("2. Background Settings")
-        handle_background_selection()
+    # 2.4: Write output video
+    log_and_progress(progress_callback, 2, 0.75, "Writing output video...")
+    try:
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height), True)
         
-        st.markdown('<div class="stage-indicator">STAGE 2: Composite with Background</div>', unsafe_allow_html=True)
+        for i, (frame, alpha) in enumerate(zip(frames, alphas)):
+            # Progress from 0.75 to 1.0 during video writing
+            progress = 0.75 + (0.25 * i / len(frames))
+            
+            if i % 30 == 0 or i == 0 or i == len(frames) - 1:
+                log_and_progress(progress_callback, 2, progress,
+                               f"Writing frame {i+1}/{len(frames)}...")
+            
+            # Create transparent output
+            alpha_3ch = np.stack([alpha] * 3, axis=-1)
+            output = (frame * alpha_3ch).astype(np.uint8)
+            out.write(output)
         
-        stage2_disabled = (not st.session_state.get('transparent_video_path') or 
-                          st.session_state.processing_stage2 or
-                          (st.session_state.bg_type == "image" and not st.session_state.get('bg_image')))
+        out.release()
+        logger.info(f"✅ Video written to {output_path}")
         
-        if st.button("🎬 Composite Final Video", type="primary", disabled=stage2_disabled,
-                    use_container_width=True, help="Combine transparent video with selected background"):
-            if st.session_state.bg_type == "image" and not st.session_state.get('bg_image'):
-                st.error("Please upload a background image first.")
-            else:
-                with st.spinner("Stage 2: Compositing with background..."):
-                    st.session_state.processing_stage2 = True
-                    try:
-                        background = st.session_state.bg_image if st.session_state.bg_type == "image" else st.session_state.bg_color
-                        final_path = stage2_composite_background(st.session_state.transparent_video_path, 
-                                                                background, st.session_state.bg_type)
-                        if final_path:
-                            st.session_state.final_video_path = final_path
-                            st.success("✅ Stage 2 Complete: Final video ready!")
-                            st.balloons()
-                        else:
-                            st.error("❌ Stage 2 Failed: Could not composite video")
-                    except Exception as e:
-                        st.error(f"❌ Stage 2 Error: {str(e)}")
-                    finally:
-                        st.session_state.processing_stage2 = False
+    except Exception as e:
+        logger.error(f"❌ Video writing failed: {e}")
+        raise
+    
+    stage2_time = time.time() - stage2_start
+    total_time = stage1_time + stage2_time
+    
+    logger.info("="*60)
+    logger.info(f"✅ STAGE 2 COMPLETE in {stage2_time:.1f}s")
+    logger.info("="*60)
+    logger.info(f"🎉 TOTAL PROCESSING TIME: {total_time:.1f}s")
+    logger.info(f"   Stage 1 (Segmentation): {stage1_time:.1f}s")
+    logger.info(f"   Stage 2 (Matting): {stage2_time:.1f}s")
+    logger.info(f"   Average: {total_time/len(frames):.2f}s per frame")
+    logger.info("="*60)
+    
+    log_and_progress(progress_callback, 2, 1.0, "Processing complete!")
+    return output_path
+
+def main():
+    st.title("🎥 Video Matting with SAM2 + MatAnyone")
+    st.write("Upload a video to remove the background from people")
+    
+    uploaded_file = st.file_uploader("Choose a video file", type=['mp4', 'avi', 'mov'])
+    
+    if uploaded_file:
+        # Create temp files
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tmp_input:
+            tmp_input.write(uploaded_file.read())
+            input_path = tmp_input.name
         
-        # Show final video result
-        if st.session_state.get('final_video_path'):
-            st.markdown("#### Final Video Result")
+        output_path = tempfile.mktemp(suffix='_output.mp4')
+        
+        if st.button("🚀 Process Video", type="primary"):
+            # Progress tracking
+            stage1_progress = st.progress(0, text="Stage 1: Initializing...")
+            stage1_status = st.empty()
+            
+            stage2_progress = st.progress(0, text="Stage 2: Waiting...")
+            stage2_status = st.empty()
+            
+            def update_progress(stage, progress, message):
+                if stage == 1:
+                    stage1_progress.progress(progress, text=f"Stage 1: {message}")
+                    stage1_status.info(f"🔄 {message}")
+                elif stage == 2:
+                    stage2_progress.progress(progress, text=f"Stage 2: {message}")
+                    stage2_status.info(f"🔄 {message}")
+            
             try:
-                with open(st.session_state.final_video_path, 'rb') as f:
-                    final_bytes = f.read()
-                st.video(final_bytes)
-                st.download_button("💾 Download Final Video (.mp4)", data=final_bytes,
-                                  file_name="final_video.mp4", mime="video/mp4", use_container_width=True)
-                st.caption(f"Size: {len(final_bytes) / (1024**2):.1f}MB")
+                logger.info("🎬 Starting video processing...")
+                result_path = process_video(input_path, output_path, update_progress)
+                
+                stage1_status.success("✅ Stage 1: Segmentation complete!")
+                stage2_status.success("✅ Stage 2: Matting complete!")
+                
+                st.success("🎉 Processing complete!")
+                
+                # Offer download
+                with open(result_path, 'rb') as f:
+                    st.download_button(
+                        label="📥 Download Result",
+                        data=f,
+                        file_name="output_matted.mp4",
+                        mime="video/mp4"
+                    )
+                
+                # Show result
+                st.video(result_path)
+                
             except Exception as e:
-                st.error(f"Error displaying final video: {str(e)}")
-        
-        # Processing tips
-        with st.expander("💡 Two-Stage Processing Tips"):
-            st.markdown("""
-            **Stage 1 - Create Transparent Video:**
-            - Uses SAM2 + MatAnyone AI to remove background
-            - Creates a .mov file with alpha channel
-            - Only needs to be done once per video
+                logger.exception("Processing failed")
+                st.error(f"❌ Processing failed: {str(e)}")
+                st.error("Check the logs above for details")
             
-            **Stage 2 - Composite Background:**
-            - Fast compositing with your chosen background
-            - Try multiple backgrounds without re-processing
-            - Much faster than Stage 1
-            """)
+            finally:
+                # Cleanup
+                for path in [input_path, output_path]:
+                    if os.path.exists(path):
+                        try:
+                            os.unlink(path)
+                        except:
+                            pass
 
 if __name__ == "__main__":
     main()