streamlit_app.py

#!/usr/bin/env python3
"""
Advanced Video Background Replacer - Streamlit Entrypoint
"""
import os
import sys
import time
from pathlib import Path
import logging
import logging.handlers
import traceback
import uuid
from tempfile import NamedTemporaryFile
import threading
import streamlit as st

from ui import render_ui
from pipeline.video_pipeline import (
    stage1_create_transparent_video,
    stage2_composite_background,
    setup_t4_environment,
    check_gpu
)
from models.model_loaders import load_sam2, load_matanyone
from utils.progress_tracker import init_progress, update_progress, mark_complete, get_progress

APP_NAME = "Advanced Video Background Replacer"
LOG_FILE = "/tmp/app.log"
LOG_MAX_BYTES = 5 * 1024 * 1024
LOG_BACKUPS = 5

def setup_logging(level: int = logging.INFO) -> logging.Logger:
    logger = logging.getLogger(APP_NAME)
    logger.setLevel(level)
    logger.propagate = False
    for h in list(logger.handlers):
        logger.removeHandler(h)
    ch = logging.StreamHandler(sys.stdout)
    ch.setLevel(level)
    ch.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
    fh = logging.handlers.RotatingFileHandler(
        LOG_FILE, maxBytes=LOG_MAX_BYTES, backupCount=LOG_BACKUPS, encoding="utf-8"
    )
    fh.setLevel(level)
    fh.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
    logger.addHandler(ch)
    logger.addHandler(fh)
    return logger

logger = setup_logging()

def custom_excepthook(type, value, tb):
    logger.error(f"Unhandled: {type.__name__}: {value}\n{''.join(traceback.format_tb(tb))}", exc_info=True)
sys.excepthook = custom_excepthook

sam2_predictor = load_sam2()
matanyone_processor = load_matanyone()

def initialize_session_state():
    defaults = {
        'uploaded_video': None,
        'video_bytes_cache': None,
        'video_preview_placeholder': None,
        'bg_image_cache': None,
        'bg_preview_placeholder': None,
        'bg_color': "#00FF00",
        'cached_color': None,
        'color_display_cache': None,
        'processed_video_bytes': None,
        'processing': False,
        'gpu_available': None,
        'last_video_id': None,
        'last_bg_image_id': None,
        'last_error': None,
        'log_level_name': 'INFO',
        'auto_refresh_logs': False,
        'log_tail_lines': 400,
        'generated_bg': None,
    }
    for k, v in defaults.items():
        if k not in st.session_state:
            st.session_state[k] = v
    if st.session_state.gpu_available is None:
        st.session_state.gpu_available = check_gpu(logger)

def process_video_background(uploaded_video, background, bg_type):
    """Background thread for video processing"""
    run_id = uuid.uuid4().hex[:8]
    logger.info("=" * 80)
    logger.info(f"[RUN {run_id}] VIDEO PROCESSING STARTED at {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())}")
    
    t0 = time.time()
    try:
        init_progress()
        update_progress("Starting video processing...", 0, "Initialization")
        
        suffix = Path(uploaded_video.name).suffix or ".mp4"
        with NamedTemporaryFile(delete=False, suffix=suffix) as tmp_vid:
            uploaded_video.seek(0)
            tmp_vid.write(uploaded_video.read())
            tmp_vid_path = tmp_vid.name
        logger.info(f"[RUN {run_id}] Temporary video path: {tmp_vid_path}")

        def progress_cb(msg):
            # Map messages to progress percentages and stages
            progress = 0
            stage = "Processing"
            
            if "Stage 1 initiated" in msg:
                progress, stage = 5, "Stage 1"
            elif "GPU engaged" in msg:
                progress, stage = 10, "Stage 1 - SAM2"
            elif "SAM2 processing frame" in msg:
                progress, stage = 15, "Stage 1 - SAM2"
            elif "SAM2 complete" in msg:
                progress, stage = 25, "Stage 1 - SAM2"
            elif "MatAnyone starting" in msg:
                progress, stage = 30, "Stage 1 - MatAnyone"
            elif "MatAnyone processing" in msg:
                progress, stage = 50, "Stage 1 - MatAnyone"
            elif "MatAnyone complete" in msg:
                progress, stage = 70, "Stage 1 - MatAnyone"
            elif "Smoothing" in msg:
                progress, stage = 75, "Stage 1 - Smoothing"
            elif "transparent video" in msg:
                progress, stage = 80, "Stage 1 - Finalizing"
            elif "Stage 1 complete" in msg:
                progress, stage = 85, "Stage 1 Complete"
            elif "Stage 2 begun" in msg:
                progress, stage = 86, "Stage 2"
            elif "Compositing" in msg:
                progress, stage = 90, "Stage 2 - Compositing"
            elif "Restoring audio" in msg:
                progress, stage = 95, "Stage 2 - Audio"
            elif "Stage 2 complete" in msg:
                progress, stage = 98, "Stage 2 Complete"
                
            update_progress(msg, progress, stage)
            logger.info(f"[PROGRESS] {msg}")

        transparent_path, audio_path = stage1_create_transparent_video(
            tmp_vid_path,
            sam2_predictor=sam2_predictor,
            matanyone_processor=matanyone_processor,
            mat_timeout_sec=300,
            progress_callback=progress_cb
        )
        
        if not transparent_path or not os.path.exists(transparent_path):
            raise RuntimeError("Stage 1 failed: Transparent video not created")
        logger.info(f"[RUN {run_id}] Stage 1 completed")

        final_path = stage2_composite_background(
            transparent_path,
            audio_path,
            background,
            bg_type.lower(),
            progress_callback=progress_cb
        )
        
        if not final_path or not os.path.exists(final_path):
            raise RuntimeError("Stage 2 failed: Final video not created")
        logger.info(f"[RUN {run_id}] Stage 2 completed")

        with open(final_path, 'rb') as f:
            st.session_state.processed_video_bytes = f.read()
        
        total = time.time() - t0
        logger.info(f"[RUN {run_id}] SUCCESS total Δ={total:.2f}s")
        mark_complete(success=True)
        
    except Exception as e:
        total = time.time() - t0
        error_msg = f"Processing Error: {str(e)} (Δ {total:.2f}s)"
        logger.error(error_msg)
        logger.error(traceback.format_exc())
        st.session_state.last_error = error_msg
        mark_complete(success=False, error=str(e))
    finally:
        st.session_state.processing = False
        logger.info(f"[RUN {run_id}] Processing finished")

def main():
    try:
        st.set_page_config(
            page_title=APP_NAME,
            page_icon="🎥",
            layout="wide",
            initial_sidebar_state="expanded"
        )
        initialize_session_state()
        render_ui(process_video_background)
    except Exception as e:
        logger.error(f"Main app error: {e}", exc_info=True)
        st.error(f"App startup failed: {str(e)}. Check logs for details.")

if __name__ == "__main__":
    setup_t4_environment()
    main()