Update streamlit_app.py

streamlit_app.py

@@ -1,65 +1,61 @@
-# ==================================================================================
-# streamlit_ui.py
-# STREAMLIT VIDEO BACKGROUND REPLACER - MAIN APPLICATION (HF-Ready Logging)
-# ==================================================================================
-# Single-button workflow: Upload video + background → Process → Download result
-# Uses SAM2 + MatAnyone pipeline with temporal smoothing
-# Adds robust logging (stdout + rotating file), in-app log tail viewer,
-# and a "Download logs" button for /tmp/app.log
-# ==================================================================================
-
-# =========================================
-# CHAPTER 1: IMPORTS AND SETUP
-# =========================================
-
-import streamlit as st
+#!/usr/bin/env python3
+"""
+Core Application Logic for Video Background Replacer
+- Handles video processing pipeline (SAM2 + MatAnyone + FFmpeg)
+- Integrates with UI (imported from ui.py)
+- Manages session state and logging
+"""
 import os
 import sys
 import time
 from pathlib import Path
-import cv2
-import numpy as np
-from PIL import Image
 import logging
 import logging.handlers
-import io
-import torch
 import traceback
 import uuid
 from datetime import datetime
 from tempfile import NamedTemporaryFile
 import subprocess
+import streamlit as st
+import cv2
+import numpy as np
+from PIL import Image
+import torch
 
-# Project Setup
-sys.path.append(str(Path(__file__).parent.absolute()))
+# Import UI components
+from ui import render_ui
 
 # Import pipeline functions
-from pipeline.video_pipeline import stage1_create_transparent_video, stage2_composite_background
+from pipeline.video_pipeline import (
+    stage1_create_transparent_video,
+    stage2_composite_background,
+    setup_t4_environment,
+    check_gpu,
+    dump_env
+)
 
+# --- Constants ---
 APP_NAME = "Advanced Video Background Replacer"
-LOG_FILE = "/tmp/app.log"  # HF Spaces: writable, survives session
+LOG_FILE = "/tmp/app.log"
 LOG_MAX_BYTES = 5 * 1024 * 1024
 LOG_BACKUPS = 5
 
-# =========================================
-# CHAPTER 2: LOGGING
-# =========================================
-
+# --- Logging Setup ---
 def setup_logging(level: int = logging.INFO) -> logging.Logger:
     logger = logging.getLogger(APP_NAME)
     logger.setLevel(level)
-    logger.propagate = False  # avoid double logs in Streamlit
+    logger.propagate = False
 
     # Clear previous handlers on rerun
     for h in list(logger.handlers):
         logger.removeHandler(h)
 
-    # Console handler (stdout → visible in HF logs)
+    # Console handler
     ch = logging.StreamHandler(sys.stdout)
     ch.setLevel(level)
     ch.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
 
-    # Rotating file handler (local tailing inside the app)
+    # Rotating file handler
     fh = logging.handlers.RotatingFileHandler(
         LOG_FILE, maxBytes=LOG_MAX_BYTES, backupCount=LOG_BACKUPS, encoding="utf-8"
     )
@@ -72,133 +68,13 @@ def setup_logging(level: int = logging.INFO) -> logging.Logger:
 
 logger = setup_logging()
 
-# Global Exception Hook
+# --- Global Exception Hook ---
 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
 
-# =========================================
-# CHAPTER 3: DIAGNOSTICS HELPERS
-# =========================================
-
-def try_run(cmd, timeout=5):
-    """Run a shell command safely and return (ok, stdout|stderr)."""
-    try:
-        out = subprocess.check_output(cmd, stderr=subprocess.STDOUT, timeout=timeout, text=True)
-        return True, out.strip()
-    except Exception as e:
-        return False, str(e)
-
-def tail_file(path: str, lines: int = 400) -> str:
-    """Read last N lines from a text file efficiently."""
-    if not os.path.exists(path):
-        return "(log file not found)"
-    try:
-        # A simple, safe tail (files are small due to rotation)
-        with open(path, "r", encoding="utf-8", errors="replace") as f:
-            content = f.readlines()
-        return "".join(content[-lines:])
-    except Exception as e:
-        return f"(failed to read log: {e})"
-
-def read_file_bytes(path: str) -> bytes:
-    """Read a file as bytes; return b'' if missing/error."""
-    try:
-        if not os.path.exists(path):
-            return b""
-        with open(path, "rb") as f:
-            return f.read()
-    except Exception:
-        return b""
-
-def check_gpu(logger: logging.Logger):
-    """Check GPU availability and log details."""
-    logger.info("=" * 60)
-    logger.info("GPU DIAGNOSTIC")
-    logger.info("=" * 60)
-
-    logger.info(f"Python: {sys.version.split()[0]}")
-    logger.info(f"Torch: {torch.__version__}")
-    logger.info(f"CUDA compiled version in torch: {torch.version.cuda}")
-
-    cuda_available = torch.cuda.is_available()
-    logger.info(f"torch.cuda.is_available(): {cuda_available}")
-
-    # Try to probe nvidia-smi if present
-    ok, out = try_run(["bash", "-lc", "command -v nvidia-smi && nvidia-smi -L && nvidia-smi -q -d MEMORY | head -n 40"], timeout=6)
-    if ok:
-        logger.info("nvidia-smi probe:\n" + out)
-    else:
-        logger.info(f"nvidia-smi probe not available: {out}")
-
-    if cuda_available:
-        try:
-            count = torch.cuda.device_count()
-            logger.info(f"CUDA Device Count: {count}")
-            for i in range(count):
-                logger.info(f"Device {i}: {torch.cuda.get_device_name(i)}")
-            torch.cuda.set_device(0)
-            logger.info("Set CUDA device 0 as default")
-            test_tensor = torch.randn(64, 64).cuda()
-            logger.info(f"GPU Test: SUCCESS on {test_tensor.device}")
-            del test_tensor
-            torch.cuda.empty_cache()
-        except Exception as e:
-            logger.error(f"GPU test failed: {e}", exc_info=True)
-    else:
-        logger.warning("CUDA NOT AVAILABLE")
-
-    logger.info("=" * 60)
-    return cuda_available
-
-def dump_env(logger: logging.Logger):
-    keys_of_interest = [
-        "HF_HOME", "HF_TOKEN", "PYTORCH_CUDA_ALLOC_CONF", "CUDA_VISIBLE_DEVICES",
-        "TORCH_ALLOW_TF32_CUBLAS_OVERRIDE", "OMP_NUM_THREADS", "NUMEXPR_MAX_THREADS",
-        "PYTHONPATH", "PATH"
-    ]
-    logger.info("=== ENV VARS (sanitized) ===")
-    for k in keys_of_interest:
-        v = os.environ.get(k)
-        if not v:
-            continue
-        # Mask tokens/secrets
-        if "TOKEN" in k or "KEY" in k or "SECRET" in k or "PASSWORD" in k:
-            v = "[REDACTED]"
-        logger.info(f"{k}={v}")
-    logger.info("============================")
-
-# =========================================
-# CHAPTER 4: STREAMLIT CONFIGURATION
-# =========================================
-
-st.set_page_config(
-    page_title=APP_NAME,
-    page_icon="🎥",
-    layout="wide",
-    initial_sidebar_state="expanded"
-)
-
-# Custom CSS
-st.markdown("""
-<style>
-    .main .block-container { padding-top: 1.0rem; 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; }
-    .small-mono { font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, "Liberation Mono", "Courier New", monospace; font-size: 12px; white-space: pre-wrap; }
-</style>
-""", unsafe_allow_html=True)
-
-# =========================================
-# CHAPTER 5: SESSION STATE
-# =========================================
+sys.excepthook = custom_excepthook
 
+# --- Session State Initialization ---
 def initialize_session_state():
     defaults = {
         'uploaded_video': None,
@@ -224,296 +100,84 @@ def initialize_session_state():
             st.session_state[k] = v
 
     if st.session_state.gpu_available is None:
-        # On first load, run diagnostics once
         dump_env(logger)
         st.session_state.gpu_available = check_gpu(logger)
 
+# --- Set Log Level ---
 def set_log_level(name: str):
     name = (name or "INFO").upper()
     lvl = getattr(logging, name, logging.INFO)
-    setup_logging(lvl)  # rebuild handlers at new level
+    setup_logging(lvl)
     global logger
     logger = logging.getLogger(APP_NAME)
     logger.setLevel(lvl)
     st.session_state.log_level_name = name
     logger.info(f"Log level set to {name}")
 
-# =========================================
-# CHAPTER 6: MAIN APP
-# =========================================
+# --- Main Processing Function ---
+def process_video(uploaded_video, bg_image, bg_color, bg_type):
+    run_id = uuid.uuid4().hex[:8]
+    logger.info("=" * 80)
+    logger.info(f"[RUN {run_id}] VIDEO PROCESSING STARTED at {datetime.utcnow().isoformat()}Z")
+    logger.info("=" * 80)
 
-def main():
-    st.title(APP_NAME)
-    st.caption("HF-ready with robust logs, live tail, and log download")
-    st.markdown("---")
+    st.session_state.processing = True
+    st.session_state.processed_video_bytes = None
+    st.session_state.last_error = None
+    t0 = time.time()
 
-    initialize_session_state()
-
-    # ------------- Sidebar: System & Logs -----------------
-    with st.sidebar:
-        st.subheader("System Status")
-
-        # Log level control
-        log_level = st.selectbox(
-            "Log level",
-            ["DEBUG", "INFO", "WARNING", "ERROR"],
-            index=["DEBUG", "INFO", "WARNING", "ERROR"].index(st.session_state.log_level_name),
-            key="log_level_select"
-        )
-        if log_level != st.session_state.log_level_name:
-            set_log_level(log_level)
-
-        if st.session_state.gpu_available:
-            try:
-                dev = torch.cuda.get_device_name(0)
-            except Exception:
-                dev = "Detected (name unavailable)"
-            st.success(f"GPU: {dev}")
-        else:
-            st.error("GPU: Not Available")
-
-        st.markdown("**Log file:** `/tmp/app.log`")
-        st.checkbox("Auto-refresh log tail (every 2s)", key="auto_refresh_logs")
-
-        # --- Download logs button ---
-        log_bytes = read_file_bytes(LOG_FILE)
-        st.download_button(
-            "Download logs",
-            data=log_bytes if log_bytes else b"Log file not available yet.",
-            file_name="app.log",
-            mime="text/plain",
-            use_container_width=True,
-            disabled=not bool(log_bytes)
-        )
-
-        st.number_input("Tail last N lines", min_value=50, max_value=5000, step=50, key="log_tail_lines")
-        if st.button("Refresh Logs"):
-            st.session_state._force_log_refresh = True  # trigger rerun
-
-        with st.expander("View Log Tail", expanded=True):
-            # Auto-refresh: trigger reruns
-            if st.session_state.auto_refresh_logs:
-                time.sleep(2)
-                st.rerun()
-            log_text = tail_file(LOG_FILE, st.session_state.log_tail_lines)
-            st.code(log_text, language="text")
-
-    # ------------- Error Display -----------------
-    if st.session_state.last_error:
-        with st.expander("⚠️ Last Error", expanded=False):
-            st.error(st.session_state.last_error)
-            if st.button("Clear Error"):
-                st.session_state.last_error = None
-                st.rerun()
-
-    col1, col2 = st.columns([1, 1], gap="large")
-
-    # ------------- Column 1: Video Upload -----------------
-    with col1:
-        st.header("1. Upload Video")
-        uploaded = st.file_uploader(
-            "Upload Video",
-            type=["mp4", "mov", "avi", "mkv", "webm"],
-            key="video_uploader"
-        )
-
-        current_video_id = id(uploaded)
-        if uploaded is not None and current_video_id != st.session_state.last_video_id:
-            logger.info(f"[UI] New video selected: name={uploaded.name}, size≈{getattr(uploaded, 'size', 'n/a')} bytes")
-            st.session_state.uploaded_video = uploaded
-            st.session_state.last_video_id = current_video_id
-            st.session_state.video_bytes_cache = None
-            st.session_state.processed_video_bytes = None
-            st.session_state.last_error = None
-
-        st.markdown("### Video Preview")
-        if st.session_state.video_preview_placeholder is None:
-            st.session_state.video_preview_placeholder = st.empty()
-
-        if st.session_state.uploaded_video is not None:
-            try:
-                if st.session_state.video_bytes_cache is None:
-                    logger.info("[UI] Caching uploaded video for preview...")
-                    st.session_state.uploaded_video.seek(0)
-                    st.session_state.video_bytes_cache = st.session_state.uploaded_video.read()
-                    logger.info(f"[UI] Cached {len(st.session_state.video_bytes_cache)/1e6:.2f} MB for preview")
-
-                with st.session_state.video_preview_placeholder.container():
-                    st.video(st.session_state.video_bytes_cache)
-
-            except Exception as e:
-                logger.error(f"[UI] Video preview error: {e}", exc_info=True)
-                st.session_state.video_preview_placeholder.error(f"Cannot display video: {e}")
-        else:
-            st.session_state.video_preview_placeholder.empty()
-
-    # ------------- Column 2: Background + Processing -----------------
-    with col2:
-        st.header("2. Background Settings")
-
-        bg_type = st.radio(
-            "Select Background Type:",
-            ["Image", "Color"],
-            horizontal=True,
-            key="bg_type_radio"
-        )
-
-        # Background Image
-        if bg_type == "Image":
-            bg_image = st.file_uploader(
-                "Upload Background Image",
-                type=["jpg", "png", "jpeg"],
-                key="bg_image_uploader"
-            )
-            current_bg_id = id(bg_image)
-            if bg_image is not None and current_bg_id != st.session_state.last_bg_image_id:
-                logger.info(f"[UI] New background image: {bg_image.name}")
-                st.session_state.last_bg_image_id = current_bg_id
-                try:
-                    st.session_state.bg_image_cache = Image.open(bg_image)
-                    # Lazy convert to display; real convert is just before Stage 2
-                except Exception as e:
-                    st.session_state.bg_image_cache = None
-                    logger.error(f"[UI] Failed to load background image: {e}", exc_info=True)
-
-            if st.session_state.bg_preview_placeholder is None:
-                st.session_state.bg_preview_placeholder = st.empty()
-
-            if st.session_state.bg_image_cache is not None:
-                with st.session_state.bg_preview_placeholder.container():
-                    st.image(st.session_state.bg_image_cache, caption="Selected Background", use_container_width=True)
-            else:
-                st.session_state.bg_preview_placeholder.empty()
-
-        # Background Color
-        else:
-            selected_color = st.color_picker(
-                "Choose Background Color",
-                st.session_state.bg_color,
-                key="color_picker"
-            )
-            if selected_color != st.session_state.cached_color:
-                logger.info(f"[UI] New background color: {selected_color}")
-                st.session_state.bg_color = selected_color
-                st.session_state.cached_color = selected_color
-
-                color_rgb = tuple(int(selected_color.lstrip('#')[i:i+2], 16) for i in (0, 2, 4))
-                color_display = np.zeros((100, 100, 3), dtype=np.uint8)
-                color_display[:, :] = color_rgb
-                st.session_state.color_display_cache = color_display
-
-            if st.session_state.bg_preview_placeholder is None:
-                st.session_state.bg_preview_placeholder = st.empty()
-
-            if st.session_state.color_display_cache is not None:
-                with st.session_state.bg_preview_placeholder.container():
-                    st.image(st.session_state.color_display_cache, caption="Selected Color", width=200)
-            else:
-                st.session_state.bg_preview_placeholder.empty()
-
-        # ------------- Processing -------------
-        st.header("3. Process Video")
-        can_process = (st.session_state.uploaded_video is not None and not st.session_state.processing)
-
-        if st.button("Process Video", disabled=not can_process, use_container_width=True):
-            run_id = uuid.uuid4().hex[:8]
-            logger.info("=" * 80)
-            logger.info(f"[RUN {run_id}] VIDEO PROCESSING STARTED at {datetime.utcnow().isoformat()}Z")
-            logger.info("=" * 80)
-
-            st.session_state.processing = True
-            st.session_state.processed_video_bytes = None
-            st.session_state.last_error = None
-
-            t0 = time.time()
-            try:
-                # Validate background
-                if bg_type == "Image":
-                    if st.session_state.bg_image_cache is None:
-                        raise RuntimeError("Background type is Image, but no image was provided.")
-                    background = st.session_state.bg_image_cache.convert("RGB")
-                    logger.info(f"[RUN {run_id}] Using IMAGE background (RGB)")
-                else:
-                    if not st.session_state.bg_color:
-                        raise RuntimeError("Background type is Color, but no color was selected.")
-                    background = st.session_state.bg_color
-                    logger.info(f"[RUN {run_id}] Using COLOR background: {background}")
-
-                # Materialize uploaded video to temp file
-                if st.session_state.video_bytes_cache is None:
-                    logger.info(f"[RUN {run_id}] Reading uploaded video into cache for processing...")
-                    st.session_state.uploaded_video.seek(0)
-                    st.session_state.video_bytes_cache = st.session_state.uploaded_video.read()
-                    logger.info(f"[RUN {run_id}] Cached {len(st.session_state.video_bytes_cache)/1e6:.2f} MB for processing")
-
-                suffix = Path(st.session_state.uploaded_video.name).suffix or ".mp4"
-                with NamedTemporaryFile(delete=False, suffix=suffix) as tmp_vid:
-                    tmp_vid.write(st.session_state.video_bytes_cache)
-                    tmp_vid_path = tmp_vid.name
-
-                # Explicit debug about Stage 1 input bytes
-                input_bytes_len = len(st.session_state.video_bytes_cache) if st.session_state.video_bytes_cache else 0
-                logger.debug(f"[RUN {run_id}] Stage 1 input bytes: {input_bytes_len}")
-
-                logger.info(f"[RUN {run_id}] Temp video path: {tmp_vid_path}")
-
-                # Stage 1
-                t1 = time.time()
-                logger.info(f"[RUN {run_id}] Stage 1: Transparent video creation START")
-                transparent_path = stage1_create_transparent_video(tmp_vid_path)
-                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: DONE → {transparent_path} (Δ {time.time()-t1:.2f}s)")
+    try:
+        # Materialize uploaded video to temp file
+        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
 
-                # Stage 2
-                t2 = time.time()
-                logger.info(f"[RUN {run_id}] Stage 2: Compositing START (bg_type={bg_type.lower()})")
-                final_path = stage2_composite_background(transparent_path, background, bg_type.lower())
-                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: DONE → {final_path} (Δ {time.time()-t2:.2f}s)")
+        # Stage 1: Create transparent video and extract audio
+        transparent_path, audio_path = stage1_create_transparent_video(tmp_vid_path)
+        if not transparent_path or not os.path.exists(transparent_path):
+            raise RuntimeError("Stage 1 failed: Transparent video not created")
 
-                # Load final into memory (Streamlit download)
-                with open(final_path, 'rb') as f:
-                    st.session_state.processed_video_bytes = f.read()
+        # Stage 2: Composite with background and restore audio
+        background = bg_image.convert("RGB") if bg_type == "Image" else bg_color
+        final_path = stage2_composite_background(transparent_path, audio_path, background, bg_type.lower())
 
-                total = time.time() - t0
-                logger.info(f"[RUN {run_id}] SUCCESS size={len(st.session_state.processed_video_bytes)/1e6:.2f}MB, total Δ={total:.2f}s")
-                st.success("Video processing complete!")
-                st.session_state.last_error = None
+        if not final_path or not os.path.exists(final_path):
+            raise RuntimeError("Stage 2 failed: Final video not created")
 
-            except Exception as e:
-                total = time.time() - t0
-                error_msg = f"[RUN {run_id}] Processing Error: {str(e)} (Δ {total:.2f}s)\n\nCheck logs for details."
-                logger.error(error_msg)
-                logger.error(traceback.format_exc())
-                st.session_state.last_error = error_msg
-                st.error(error_msg)
+        # Load final video into memory for download
+        with open(final_path, 'rb') as f:
+            st.session_state.processed_video_bytes = f.read()
 
-            finally:
-                st.session_state.processing = False
-                logger.info(f"[RUN {run_id}] Processing finished (processing flag cleared)")
+        total = time.time() - t0
+        logger.info(f"[RUN {run_id}] SUCCESS size={len(st.session_state.processed_video_bytes)/1e6:.2f}MB, total Δ={total:.2f}s")
+        return True
 
-        # Results
-        if st.session_state.processed_video_bytes is not None:
-            st.markdown("---")
-            st.markdown("### Processed Video")
-            try:
-                st.video(st.session_state.processed_video_bytes)
-                st.download_button(
-                    label="Download Processed Video",
-                    data=st.session_state.processed_video_bytes,
-                    file_name="processed_video.mp4",
-                    mime="video/mp4",
-                    use_container_width=True
-                )
-            except Exception as e:
-                logger.error(f"[UI] Display error: {e}", exc_info=True)
-                st.error(f"Display error: {e}")
+    except Exception as e:
+        total = time.time() - t0
+        error_msg = f"[RUN {run_id}] Processing Error: {str(e)} (Δ {total:.2f}s)\n\nCheck logs for details."
+        logger.error(error_msg)
+        logger.error(traceback.format_exc())
+        st.session_state.last_error = error_msg
+        return False
+
+    finally:
+        st.session_state.processing = False
+        logger.info(f"[RUN {run_id}] Processing finished")
+
+# --- Main App Entry Point ---
+def main():
+    st.set_page_config(
+        page_title=APP_NAME,
+        page_icon="🎥",
+        layout="wide",
+        initial_sidebar_state="expanded"
+    )
 
-# =========================================
-# CHAPTER 7: ENTRY POINT
-# =========================================
+    initialize_session_state()
+    render_ui(process_video)
 
 if __name__ == "__main__":
+    setup_t4_environment()
     main()