virtual_bg_pipeline.py

#!/usr/bin/env python3
"""
WhatsApp Virtual Background - 3D Room Compositing Pipeline
Uses MediaPipe for real-time segmentation and OpenCV for compositing
"""

import cv2
import numpy as np
import mediapipe as mp
from typing import Optional, Tuple
import warnings

warnings.filterwarnings('ignore')

class VirtualBackgroundEngine:
    def __init__(self, bg_image_path: Optional[str] = None, bg_3d_path: Optional[str] = None):
        """
        Initialize the virtual background engine

        Args:
            bg_image_path: Path to the 2D background image
            bg_3d_path: Path to the 3D scene file (GLB/GLTF/OBJ/FBX)
        """
        self.mp_selfie = mp.solutions.selfie_segmentation
        self.segmentation = self.mp_selfie.SelfieSegmentation(model_selection=1)

        self.background_2d = None
        self.background_3d = None
        self.depth_map = None

        if bg_image_path:
            self.load_2d_background(bg_image_path)
        elif bg_3d_path:
            self.load_3d_background(bg_3d_path)
        else:
            self.background_2d = self._create_default_background()

    def load_2d_background(self, image_path: str):
        """Load a static 2D image as background"""
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError(f"Could not load image: {image_path}")
        # Convert BGR to RGB
        self.background_2d = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        print(f"Loaded 2D background: {self.background_2d.shape}")

    def load_3d_background(self, path_3d: str):
        """
        Load 3D scene file (GLB, GLTF, OBJ, FBX)
        For now we extract the default view as a 2D background
        """
        # TODO: Implement full 3D rendering with PyOpenGL/Panda3D
        # For MVP, we will render the 3D scene to a 2D texture
        try:
            import trimesh
            scene = trimesh.load(path_3d)
            if isinstance(scene, trimesh.Scene):
                # Render scene to image
                self.background_2d = scene.save_image(resolution=(1920, 1080))
            else:
                # Single mesh
                self.background_2d = scene.save_image(resolution=(1920, 1080))
            print(f"Loaded 3D background from: {path_3d}")
        except ImportError:
            print("trimesh not installed, falling back to 2D background")
            self.background_2d = self._create_default_background()

    def _create_default_background(self) -> np.ndarray:
        """Create a default gradient background"""
        bg = np.zeros((1080, 1920, 3), dtype=np.uint8)
        # Create a subtle gradient
        for i in range(1080):
            color = int(50 + (i / 1080) * 40)
            bg[i, :] = [color, color, color + 10]
        return bg

    def _resize_background(self, frame: np.ndarray) -> np.ndarray:
        """Resize background to match frame dimensions"""
        h, w = frame.shape[:2]
        if self.background_2d is not None:
            return cv2.resize(self.background_2d, (w, h))
        return self._create_default_background()[:h, :w]

    def apply(self, frame: np.ndarray) -> np.ndarray:
        """
        Apply virtual background to a frame

        Args:
            frame: Input BGR frame from webcam

        Returns:
            Composited frame with virtual background
        """
        # Convert to RGB for MediaPipe
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        h, w = frame.shape[:2]

        # Get segmentation mask
        results = self.segmentation.process(frame_rgb)
        mask = results.segmentation_mask

        # Threshold and blur mask for smoother edges
        mask = (mask > 0.5).astype(np.uint8) * 255
        mask = cv2.GaussianBlur(mask, (7, 7), 0)
        mask = mask.astype(np.float32) / 255.0

        # Resize mask to frame size
        mask = cv2.resize(mask, (w, h))
        mask = np.expand_dims(mask, axis=-1)

        # Get background
        bg = self._resize_background(frame)

        # Composite: foreground * mask + background * (1 - mask)
        composited = (frame_rgb * mask + bg * (1 - mask)).astype(np.uint8)

        # Convert back to BGR for OpenCV display
        return cv2.cvtColor(composited, cv2.COLOR_RGB2BGR)

    def apply_with_depth(self, frame: np.ndarray, depth_map: np.ndarray) -> np.ndarray:
        """
        Apply background with depth-aware compositing
        (for 3D scenes where you want depth ordering)

        Args:
            frame: Input BGR frame
            depth_map: Depth map (grayscale, 0-255, higher = closer)

        Returns:
            Depth-aware composited frame
        """
        # TODO: Implement depth-aware compositing
        # For now, fall back to standard compositing
        return self.apply(frame)

    def release(self):
        """Clean up resources"""
        self.segmentation.close()


def main():
    """Main loop for testing virtual background"""
    # Initialize
    engine = VirtualBackgroundEngine(bg_image_path="background.jpg")

    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        print("Error: Could not open webcam")
        return

    print("Press 'q' to quit, 's' to save screenshot")

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        # Apply virtual background
        result = engine.apply(frame)

        # Display
        cv2.imshow('Virtual Background', result)
        # Also show original for comparison
        cv2.imshow('Original', frame)

        key = cv2.waitKey(1) & 0xFF
        if key == ord('q'):
            break
        elif key == ord('s'):
            cv2.imwrite('virtual_bg_screenshot.png', result)
            print("Screenshot saved!")

    cap.release()
    cv2.destroyAllWindows()
    engine.release()


if __name__ == "__main__":
    main()