depth_estimator.py

#!/usr/bin/env python3
"""
Depth estimation module for realistic 3D compositing
Uses MiDaS for monocular depth estimation
"""

import cv2
import numpy as np
import torch
import warnings
from typing import Optional

warnings.filterwarnings('ignore')


class DepthEstimator:
    """
    Depth estimation using MiDaS for realistic perspective matching
    """

    def __init__(self, model_type: str = "DPT_Large"):
        """
        Initialize depth estimator

        Args:
            model_type: MiDaS model variant ("DPT_Large", "DPT_Hybrid", "MiDaS_small")
        """
        self.model_type = model_type
        self.model = None
        self.device = None
        self.transform = None

        try:
            # Try to load MiDaS model
            self.model = torch.hub.load("intel-isl/MiDaS", model_type)
            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            self.model.to(self.device)
            self.model.eval()

            # Load transforms
            midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
            self.transform = midas_transforms.dpt_transform if "DPT" in model_type else midas_transforms.small_transform

            print(f"Depth estimator loaded: {model_type} on {self.device}")
        except Exception as e:
            print(f"Could not load MiDaS: {e}")
            print("Falling back to simple depth estimation")

    def estimate(self, frame: np.ndarray) -> np.ndarray:
        """
        Estimate depth map from a frame

        Args:
            frame: BGR image

        Returns:
            Depth map (normalized 0-255)
        """
        if self.model is None:
            return self._simple_depth(frame)

        try:
            # Preprocess
            img_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            input_batch = self.transform(img_rgb).to(self.device)

            # Predict depth
            with torch.no_grad():
                prediction = self.model(input_batch)
                prediction = torch.nn.functional.interpolate(
                    prediction.unsqueeze(1),
                    size=frame.shape[:2],
                    mode="bicubic",
                    align_corners=False,
                ).squeeze()

            # Convert to numpy and normalize
            depth = prediction.cpu().numpy()
            depth = cv2.normalize(depth, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)

            return depth
        except Exception as e:
            print(f"Depth estimation failed: {e}")
            return self._simple_depth(frame)

    def _simple_depth(self, frame: np.ndarray) -> np.ndarray:
        """
        Fallback simple depth estimation based on focal distance
        (Objects higher in frame = further away, roughly)
        """
        h, w = frame.shape[:2]
        # Simple gradient: closer at bottom, further at top
        depth = np.linspace(128, 255, h, dtype=np.uint8)
        depth = np.tile(depth[:, np.newaxis], (1, w))
        return depth

    def match_perspective(self, person_depth: np.ndarray, bg_depth: np.ndarray, 
                          threshold: float = 0.3) -> np.ndarray:
        """
        Match person perspective to background depth

        Args:
            person_depth: Depth map of person
            bg_depth: Depth map of background scene
            threshold: Depth matching threshold

        Returns:
            Adjusted depth mask for compositing
        """
        # Normalize both to 0-1
        p_norm = person_depth / 255.0
        b_norm = bg_depth / 255.0

        # Determine if person is in front or behind background objects
        depth_diff = p_norm - b_norm

        # Create mask: 1 where person is closer, 0 where background is closer
        mask = (depth_diff < threshold).astype(np.float32)

        return mask


def load_3d_depth_map(depth_image_path: str) -> np.ndarray:
    """
    Load a pre-rendered depth map from a 3D scene export

    Args:
        depth_image_path: Path to depth map image (grayscale)

    Returns:
        Depth map as numpy array
    """
    depth = cv2.imread(depth_image_path, cv2.IMREAD_GRAYSCALE)
    if depth is None:
        raise ValueError(f"Could not load depth map: {depth_image_path}")
    return depth


if __name__ == "__main__":
    # Test depth estimation
    import time

    print("Testing depth estimator...")
    estimator = DepthEstimator()

    # Create test frame
    test_frame = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)

    start = time.time()
    depth = estimator.estimate(test_frame)
    elapsed = time.time() - start

    print(f"Depth estimation took {elapsed:.3f}s")
    print(f"Depth map shape: {depth.shape}, range: [{depth.min()}, {depth.max()}]")