lighting_matcher.py

#!/usr/bin/env python3
"""
Lighting and Color Matching Module
Ensures realistic integration of person into virtual 3D space
"""

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

warnings.filterwarnings('ignore')


class LightingMatcher:
    """
    Matches lighting and color of foreground to background
    for realistic compositing
    """

    def __init__(self):
        self.prev_foreground = None
        self.adaptation_rate = 0.1

    def match_lighting(self, foreground: np.ndarray, background: np.ndarray,
                       mask: np.ndarray) -> np.ndarray:
        """
        Match lighting of foreground to background

        Args:
            foreground: Person RGB image
            background: Background RGB image
            mask: Segmentation mask (0-1 float)

        Returns:
            Lighting-adjusted foreground
        """
        # Calculate mean color stats for background
        bg_mean = np.mean(background, axis=(0, 1))
        fg_mean = np.mean(foreground, axis=(0, 1))

        # Simple color correction to match ambient light
        color_correction = bg_mean / (fg_mean + 1e-6)
        color_correction = np.clip(color_correction, 0.5, 2.0)

        # Apply gentle correction
        adjusted = foreground.copy().astype(np.float32)
        for i in range(3):
            adjusted[:, :, i] *= color_correction[i] * 0.3 + 0.7  # Blend with original

        # Apply shadow estimation
        shadow = self._estimate_shadow(background, mask)
        adjusted = (adjusted * (1 - shadow * 0.3)).astype(np.uint8)

        return np.clip(adjusted, 0, 255).astype(np.uint8)

    def match_color_temperature(self, foreground: np.ndarray,
                                background: np.ndarray) -> np.ndarray:
        """
        Match color temperature (warmth/coolness)

        Args:
            foreground: Person RGB image
            background: Background RGB image

        Returns:
            Color-temperature-matched foreground
        """
        # Convert to LAB color space for better color matching
        fg_lab = cv2.cvtColor(foreground, cv2.COLOR_RGB2LAB)
        bg_lab = cv2.cvtColor(background, cv2.COLOR_RGB2LAB)

        # Match a-channel (green-red) and b-channel (blue-yellow)
        for i in range(1, 3):
            fg_mean = np.mean(fg_lab[:, :, i])
            bg_mean = np.mean(bg_lab[:, :, i])
            fg_std = np.std(fg_lab[:, :, i])
            bg_std = np.std(bg_lab[:, :, i])

            # Match mean and standard deviation
            fg_lab[:, :, i] = (fg_lab[:, :, i] - fg_mean) * (bg_std / (fg_std + 1e-6)) + bg_mean

        # Convert back to RGB
        matched = cv2.cvtColor(fg_lab, cv2.COLOR_LAB2RGB)
        return matched

    def _estimate_shadow(self, background: np.ndarray, mask: np.ndarray) -> np.ndarray:
        """
        Estimate where shadows should fall based on background geometry
        """
        # Simple shadow estimation based on background darkness
        bg_gray = cv2.cvtColor(background, cv2.COLOR_RGB2GRAY)
        bg_darkness = 255 - bg_gray

        # Blur to create soft shadows
        shadow = cv2.GaussianBlur(bg_darkness, (31, 31), 0)
        shadow = shadow / 255.0 * 0.5  # Scale shadow intensity

        return shadow.astype(np.float32)

    def add_ambient_occlusion(self, image: np.ndarray, depth_map: np.ndarray) -> np.ndarray:
        """
        Add ambient occlusion based on depth map
        (darker areas where surfaces meet)

        Args:
            image: RGB image
            depth_map: Depth map (0-255, higher = closer)

        Returns:
            Image with AO applied
        """
        # Simple edge detection on depth map for AO
        edges = cv2.Canny(depth_map, 50, 150)
        ao = cv2.GaussianBlur(edges, (15, 15), 0)
        ao = ao.astype(np.float32) / 255.0

        # Apply AO to darken edges
        result = image.astype(np.float32) * (1 - ao[:, :, np.newaxis] * 0.3)
        return result.astype(np.uint8)

    def blend_edges(self, foreground: np.ndarray, background: np.ndarray,
                    mask: np.ndarray, edge_width: int = 5) -> np.ndarray:
        """
        Blend edges of person with background for seamless integration

        Args:
            foreground: Person RGB image
            background: Background RGB image
            mask: Segmentation mask (0-1 float)
            edge_width: Width of edge blending zone

        Returns:
            Blended image
        """
        # Create distance transform for edge weighting
        mask_8u = (mask * 255).astype(np.uint8)
        dist = cv2.distanceTransform(255 - mask_8u, cv2.DIST_L2, 5)

        # Normalize to edge width
        blend_weight = np.clip(dist / edge_width, 0, 1)
        blend_weight = np.expand_dims(blend_weight, axis=-1)

        # Blend
        result = foreground * mask + background * (1 - mask)
        edge_blend = foreground * blend_weight + background * (1 - blend_weight)

        # Apply edge blend only where mask is transitioning
        result = np.where(mask > 0.1, edge_blend, result)

        return result.astype(np.uint8)

    def apply_vignette(self, image: np.ndarray, strength: float = 0.3) -> np.ndarray:
        """
        Apply subtle vignette to match typical video call aesthetics

        Args:
            image: RGB image
            strength: Vignette strength (0-1)

        Returns:
            Image with vignette
        """
        rows, cols = image.shape[:2]

        # Create vignette mask
        X_resultant = cv2.getGaussianKernel(cols, cols/2)
        Y_resultant = cv2.getGaussianKernel(rows, rows/2)

        kernel = X_resultant * Y_resultant.T
        mask = 255 * kernel / np.linalg.norm(kernel)

        # Apply mask
        mask = np.dstack([mask] * 3)
        result = image.astype(np.float32) * (0.7 + mask * 0.3 / 255)

        return np.clip(result, 0, 255).astype(np.uint8)


class PostProcessor:
    """
    Final post-processing for the composited image
    """

    @staticmethod
    def denoise(image: np.ndarray) -> np.ndarray:
        """Apply light denoising"""
        return cv2.fastNlMeansDenoisingColored(image, None, 5, 5, 7, 21)

    @staticmethod
    def enhance_details(image: np.ndarray) -> np.ndarray:
        """Enhance fine details"""
        kernel = np.array([[-1, -1, -1],
                          [-1,  9, -1],
                          [-1, -1, -1]])
        return cv2.filter2D(image, -1, kernel)

    @staticmethod
    def adjust_exposure(image: np.ndarray, gamma: float = 1.0) -> np.ndarray:
        """Adjust exposure with gamma correction"""
        lookup_table = np.array([((i / 255.0) ** gamma) * 255
                                for i in range(256)]).astype(np.uint8)
        return cv2.LUT(image, lookup_table)


if __name__ == "__main__":
    # Test lighting matcher
    print("Testing Lighting Matcher...")

    # Create test images
    fg = np.ones((720, 1280, 3), dtype=np.uint8) * 200
    bg = np.ones((720, 1280, 3), dtype=np.uint8) * 100
    mask = np.zeros((720, 1280, 1), dtype=np.float32)
    mask[200:500, 400:800] = 1.0

    matcher = LightingMatcher()
    result = matcher.match_lighting(fg, bg, mask)
    result = matcher.match_color_temperature(result, bg)

    print(f"Input: {fg.shape}, Output: {result.shape}")
    print("Lighting matcher test complete!")