utils/blur_engine.py

"""
Blur Engine - Applies various blur effects using segmentation masks
Supports Gaussian blur, pixelation, and solid color overlay
"""
import cv2
import numpy as np
from typing import Optional


class BlurEngine:
    """Applies blur effects to images using segmentation masks"""
    
    def __init__(self, config):
        self.config = config
    
    def apply_blur(
        self,
        frame: np.ndarray,
        mask: np.ndarray,
        blur_type: Optional[str] = None,
        blur_strength: Optional[int] = None,
        edge_feather: Optional[int] = None,
    ) -> np.ndarray:
        """
        Apply blur to masked region of frame
        
        Args:
            frame: BGR image (H, W, 3) uint8
            mask: Binary mask (H, W) uint8, 255 = blur region
            blur_type: "gaussian", "pixelate", or "black"
            blur_strength: Kernel size for Gaussian (must be odd)
            edge_feather: Edge softening kernel (must be odd)
        
        Returns:
            Blurred frame (H, W, 3) uint8
        """
        blur_type = blur_type or self.config.blur.blur_type
        blur_strength = blur_strength or self.config.blur.blur_strength
        edge_feather = edge_feather or self.config.blur.edge_feather
        
        # Ensure odd kernel sizes
        blur_strength = blur_strength | 1
        edge_feather = edge_feather | 1
        
        # Skip if no mask
        if mask is None or mask.max() == 0:
            return frame.copy()
        
        # Resize mask to match frame if needed
        if mask.shape[:2] != frame.shape[:2]:
            mask = cv2.resize(mask, (frame.shape[1], frame.shape[0]), 
                            interpolation=cv2.INTER_LINEAR)
            mask = (mask > 127).astype(np.uint8) * 255
        
        # Generate the blurred/modified version
        if blur_type == "gaussian":
            modified = self._gaussian_blur(frame, blur_strength)
        elif blur_type == "pixelate":
            modified = self._pixelate(frame, self.config.blur.pixelate_size)
        elif blur_type == "black":
            modified = np.zeros_like(frame)
        else:
            modified = self._gaussian_blur(frame, blur_strength)
        
        # Apply edge feathering for smooth transitions
        mask_soft = self._feather_mask(mask, edge_feather)
        
        # Composite: blend original and modified using soft mask
        result = self._composite(frame, modified, mask_soft)
        
        return result
    
    def _gaussian_blur(self, frame: np.ndarray, kernel_size: int) -> np.ndarray:
        """Apply Gaussian blur to entire frame"""
        return cv2.GaussianBlur(frame, (kernel_size, kernel_size), 0)
    
    def _pixelate(self, frame: np.ndarray, block_size: int) -> np.ndarray:
        """Apply pixelation effect"""
        h, w = frame.shape[:2]
        small = cv2.resize(frame, (w // block_size, h // block_size), 
                          interpolation=cv2.INTER_LINEAR)
        pixelated = cv2.resize(small, (w, h), interpolation=cv2.INTER_NEAREST)
        return pixelated
    
    def _feather_mask(self, mask: np.ndarray, kernel_size: int) -> np.ndarray:
        """
        Create soft-edged mask for smooth blur transitions
        
        Returns:
            Float mask (H, W) in range [0, 1]
        """
        # Optional: dilate slightly to ensure full coverage
        dilate_kernel = np.ones((3, 3), np.uint8)
        mask_dilated = cv2.dilate(mask, dilate_kernel, iterations=1)
        
        # Apply Gaussian blur to create soft edges
        mask_soft = cv2.GaussianBlur(
            mask_dilated.astype(np.float32), 
            (kernel_size, kernel_size), 
            0
        )
        
        # Normalize to [0, 1]
        mask_max = mask_soft.max()
        if mask_max > 0:
            mask_soft = mask_soft / mask_max
        
        return mask_soft
    
    def _composite(
        self, 
        original: np.ndarray, 
        modified: np.ndarray, 
        mask_soft: np.ndarray
    ) -> np.ndarray:
        """
        Alpha-blend original and modified frames using soft mask
        
        result = original * (1 - mask) + modified * mask
        """
        # Expand mask to 3 channels
        mask_3ch = np.stack([mask_soft] * 3, axis=-1)
        
        # Blend
        result = (
            original.astype(np.float32) * (1.0 - mask_3ch) + 
            modified.astype(np.float32) * mask_3ch
        )
        
        return np.clip(result, 0, 255).astype(np.uint8)
    
    def visualize_mask(
        self, 
        frame: np.ndarray, 
        mask: np.ndarray, 
        color: tuple = (0, 255, 0),
        alpha: float = 0.4
    ) -> np.ndarray:
        """
        Overlay colored mask on frame for visualization
        
        Args:
            frame: BGR image
            mask: Binary mask (H, W) uint8
            color: BGR color tuple
            alpha: Overlay transparency
        
        Returns:
            Visualization image
        """
        vis = frame.copy()
        
        if mask is None or mask.max() == 0:
            return vis
        
        # Resize mask if needed
        if mask.shape[:2] != frame.shape[:2]:
            mask = cv2.resize(mask, (frame.shape[1], frame.shape[0]))
        
        # Create colored overlay
        overlay = np.zeros_like(frame)
        overlay[mask > 127] = color
        
        # Blend
        vis = cv2.addWeighted(vis, 1.0, overlay, alpha, 0)
        
        # Draw contours
        contours, _ = cv2.findContours(
            mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )
        cv2.drawContours(vis, contours, -1, color, 2)
        
        return vis