moge/utils/data_augmentation.py

import os
import json
import time
import random
from typing import *
import itertools
from numbers import Number
import io

import numpy as np
import cv2
from PIL import Image
import torch
import torchvision.transforms.v2.functional as TF
try:
    import utils3d_moge as utils3d
except ImportError:
    import utils3d
from scipy.signal import fftconvolve

from ..utils.geometry_numpy import harmonic_mean_numpy, norm3d, depth_occlusion_edge_numpy


def sample_perspective(
    src_intrinsics: np.ndarray, 
    tgt_aspect: float, 
    center_augmentation: float, 
    fov_range_absolute: Tuple[float, float], 
    fov_range_relative: Tuple[float, float], 
    rng: np.random.Generator = None
) -> Tuple[np.ndarray, np.ndarray]:  
    raw_horizontal, raw_vertical = abs(1.0 / src_intrinsics[0, 0]), abs(1.0 / src_intrinsics[1, 1])
    raw_fov_x, raw_fov_y = utils3d.np.intrinsics_to_fov(src_intrinsics)

    # 1. set target fov
    fov_range_absolute_min, fov_range_absolute_max = fov_range_absolute
    fov_range_relative_min, fov_range_relative_max = fov_range_relative
    tgt_fov_x_min = min(fov_range_relative_min * raw_fov_x, utils3d.focal_to_fov(utils3d.fov_to_focal(fov_range_relative_min * raw_fov_y) / tgt_aspect))
    tgt_fov_x_max = min(fov_range_relative_max * raw_fov_x, utils3d.focal_to_fov(utils3d.fov_to_focal(fov_range_relative_max * raw_fov_y) / tgt_aspect))
    tgt_fov_x_min, tgt_fov_max = max(np.deg2rad(fov_range_absolute_min), tgt_fov_x_min), min(np.deg2rad(fov_range_absolute_max), tgt_fov_x_max)
    tgt_fov_x = rng.uniform(min(tgt_fov_x_min, tgt_fov_x_max), tgt_fov_x_max)
    tgt_fov_y = utils3d.focal_to_fov(utils3d.np.fov_to_focal(tgt_fov_x) * tgt_aspect)

    # 2. set target image center (principal point) and the corresponding z-direction in raw camera space
    center_dtheta = center_augmentation * rng.uniform(-0.5, 0.5) * (raw_fov_x - tgt_fov_x)
    center_dphi = center_augmentation * rng.uniform(-0.5, 0.5) * (raw_fov_y - tgt_fov_y)
    cu, cv = 0.5 + 0.5 * np.tan(center_dtheta) / np.tan(raw_fov_x / 2), 0.5 + 0.5 *  np.tan(center_dphi) / np.tan(raw_fov_y / 2)
    direction = utils3d.np.unproject_cv(np.array([[cu, cv]], dtype=np.float32), np.array([1.0], dtype=np.float32), intrinsics=src_intrinsics)[0]
    
    # 3. obtain the rotation matrix for homography warping (new_ext = R * old_ext)
    R = utils3d.np.rotation_matrix_from_vectors(direction, np.array([0, 0, 1], dtype=np.float32))

    # 4. shrink the target view to fit into the warped image
    corners = np.array([[0, 0], [0, 1], [1, 1], [1, 0]], dtype=np.float32)
    corners = np.concatenate([corners, np.ones((4, 1), dtype=np.float32)], axis=1) @ (np.linalg.inv(src_intrinsics).T @ R.T)   # corners in viewport's camera plane
    corners = corners[:, :2] / corners[:, 2:3]
    tgt_horizontal, tgt_vertical = np.tan(tgt_fov_x / 2) * 2, np.tan(tgt_fov_y / 2) * 2
    warp_horizontal, warp_vertical = float('inf'), float('inf')
    for i in range(4):
        intersection, _ = utils3d.np.ray_intersection(
            np.array([0., 0.]), np.array([[tgt_aspect, 1.0], [tgt_aspect, -1.0]]),
            corners[i - 1], corners[i] - corners[i - 1],
        )
        warp_horizontal, warp_vertical = min(warp_horizontal, 2 * np.abs(intersection[:, 0]).min()), min(warp_vertical, 2 * np.abs(intersection[:, 1]).min())
    tgt_horizontal, tgt_vertical = min(tgt_horizontal, warp_horizontal), min(tgt_vertical, warp_vertical)
    
    # 5. obtain the target intrinsics
    fx, fy = 1 / tgt_horizontal, 1 / tgt_vertical
    tgt_intrinsics = utils3d.np.intrinsics_from_focal_center(fx, fy, 0.5, 0.5).astype(np.float32)

    return tgt_intrinsics, R


def warp_perspective(
    src_map: np.ndarray = None, 
    transform: np.ndarray = None,
    tgt_size: Tuple[int, int] = None, 
    interpolation: Literal['nearest', 'bilinear', 'lanczos'] = 'nearest',
    sparse_mask: np.ndarray = None,
):
    """Perspective warping with careful resampling.
    - For `lanczos`, use PIL to resize first to reduce aliasing.
    - For `nearest` with sparse input, use mask-aware nearest resize to avoid losing points.
    - For `bilinear` or `nearest` with dense input, directly use cv2.remap.

    - `transform` is the matrix that transforms homogeneous pixel coordinates of source image to those of target image, i.e., `p_tgt = transform @ p_src`.
    """
    
    tgt_height, tgt_width = tgt_size
    src_height, src_width = src_map.shape[:2]

    # source to target transform
    transform_pixel = np.array([[tgt_width, 0, -0.5], [0, tgt_height, -0.5], [0, 0, 1]], dtype=np.float32) @ transform @ np.array([[1 / src_width, 0, 0.5 / src_width], [0, 1 / src_height, 0.5 / src_height], [0, 0, 1]], dtype=np.float32)
    # Get scale factor at the target center
    w = np.dot(np.linalg.inv(transform_pixel)[2, :], np.array([tgt_width / 2, tgt_height / 2, 1], dtype=np.float32))
    scale_x, scale_y = w * np.linalg.norm(transform_pixel[:2, :2], axis=0)

    if interpolation == 'lanczos' and (scale_x < 0.8 or scale_y < 0.8):
        # If lanczos & downsampling, use PIL to resize first to reduce aliasing
        src_height, src_width = max(round(src_height * scale_y * 1.25), 16), max(round(src_width * scale_x * 1.25), 16)
        src_map = np.array(Image.fromarray(src_map).resize((src_width, src_height), Image.Resampling.LANCZOS))
    elif interpolation == 'nearest' and sparse_mask is not None and (scale_x < 1 or scale_y < 1):
        # If nearest and sparse, use mask-aware nearest resize first to avoid losing points
        src_height, src_width = max(round(src_height * scale_y), 16), max(round(src_width * scale_x), 16)
        src_map, _ = utils3d.np.masked_nearest_resize(src_map, mask=sparse_mask, size=(src_height, src_width))

    # Recompute the pixel-space transform after resizing
    transform_pixel = np.array([[tgt_width, 0, -0.5], [0, tgt_height, -0.5], [0, 0, 1]], dtype=np.float32) @ transform @ np.array([[1 / src_width, 0, 0.5 / src_width], [0, 1 / src_height, 0.5 / src_height], [0, 0, 1]], dtype=np.float32)
    
    # Remap
    cv2_interpolation = {'nearest': cv2.INTER_NEAREST, 'bilinear': cv2.INTER_LINEAR, 'lanczos': cv2.INTER_LANCZOS4}[interpolation]
    tgt_map = cv2.warpPerspective(src_map, transform_pixel, (tgt_width, tgt_height), flags=cv2_interpolation)

    return tgt_map


def image_color_augmentation(image: np.ndarray, augmentations: List[Dict[str, Any]], rng: np.random.Generator = None, depth: np.ndarray = None):
    height, width = image.shape[:2]
    if rng is None:
        rng = np.random.default_rng()
    if 'jittering' in augmentations:
        image = torch.from_numpy(image).permute(2, 0, 1)
        image = TF.adjust_brightness(image, rng.uniform(0.9, 1.1))
        image = TF.adjust_contrast(image, rng.uniform(0.9, 1.1))
        image = TF.adjust_saturation(image, rng.uniform(0.9, 1.1))
        image = TF.adjust_hue(image, rng.uniform(-0.05, 0.05))
        image = TF.adjust_gamma(image, rng.uniform(0.9, 1.1))
        image = image.permute(1, 2, 0).numpy()
    if 'dof' in augmentations:
        assert depth is not None, 'Depth map is required for DOF augmentation'
        if rng.uniform() < 0.5:
            dof_strength = rng.integers(12)
            disp = 1 / depth
            finite_mask = np.isfinite(depth)
            disp_min, disp_max = disp[finite_mask].min(), disp[finite_mask].max()
            disp = cv2.inpaint(np.nan_to_num(disp, nan=1), np.isnan(disp).astype(np.uint8), 3, cv2.INPAINT_TELEA).clip(0, disp_max)
            dof_focus = rng.uniform(disp_min, disp_max)
            image = depth_of_field(image, disp, dof_focus, dof_strength)
    if 'shot_noise' in augmentations:
        if rng.uniform() < 0.5: 
            k = np.exp(rng.uniform(np.log(100), np.log(10000))) / 255
            image = (rng.poisson(image * k) / k).clip(0, 255).astype(np.uint8)
    if 'blurring' in augmentations:
        if rng.uniform() < 0.5:    
            ratio = rng.uniform(0.25, 1)
            image = cv2.resize(cv2.resize(image, (int(width * ratio), int(height * ratio)), interpolation=cv2.INTER_AREA), (width, height), interpolation=rng.choice([cv2.INTER_LINEAR_EXACT, cv2.INTER_CUBIC, cv2.INTER_LANCZOS4]))
    if 'jpeg_loss' in augmentations:
        if rng.uniform() < 0.5: 
            image = cv2.imdecode(cv2.imencode('.jpg', image, [cv2.IMWRITE_JPEG_QUALITY, rng.integers(20, 100)])[1], cv2.IMREAD_COLOR)
    
    return image



def disk_kernel(radius: int) -> np.ndarray:
    """
    Generate disk kernel with given radius.
    
    Args:
        radius (int): Radius of the disk (in pixels).
    
    Returns:
        np.ndarray: (2*radius+1, 2*radius+1) normalized convolution kernel.
    """
    # Create coordinate grid centered at (0,0)
    L = np.arange(-radius, radius + 1)
    X, Y = np.meshgrid(L, L)
    # Generate disk: region inside circle with radius R is 1
    kernel = ((X**2 + Y**2) <= radius**2).astype(np.float32)
    # Normalize the kernel
    kernel /= np.sum(kernel)
    return kernel


def disk_blur(image: np.ndarray, radius: int) -> np.ndarray:
    """
    Apply disk blur to an image using FFT convolution.

    Args:
        image (np.ndarray): Input image, can be grayscale or color.
        radius (int): Blur radius (in pixels).

    Returns:
        np.ndarray: Blurred image.
    """
    if radius == 0:
        return image
    kernel = disk_kernel(radius)
    if image.ndim == 2:
        blurred = fftconvolve(image, kernel, mode='same')
    elif image.ndim == 3:
        channels = []
        for i in range(image.shape[2]):
            blurred_channel = fftconvolve(image[..., i], kernel, mode='same')
            channels.append(blurred_channel)
        blurred = np.stack(channels, axis=-1)
    else:
        raise ValueError("Image must be 2D or 3D.")
    return blurred


def depth_of_field(
    img: np.ndarray, 
    disp: np.ndarray, 
    focus_disp : float, 
    max_blur_radius : int = 10,
) -> np.ndarray:
    """
    Apply depth of field effect to an image.

    Args:
        img (numpy.ndarray): (H, W, 3) input image.
        depth (numpy.ndarray): (H, W) depth map of the scene.
        focus_depth (float): Focus depth of the lens.
        strength (float): Strength of the depth of field effect.
        max_blur_radius (int): Maximum blur radius (in pixels).
        
    Returns:
        numpy.ndarray: (H, W, 3) output image with depth of field effect applied.
    """
    # Precalculate dialated depth map for each blur radius
    max_disp = np.max(disp)
    disp = disp / max_disp
    focus_disp = focus_disp / max_disp
    dilated_disp = []
    for radius in range(max_blur_radius + 1):
        dilated_disp.append(cv2.dilate(disp, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * radius + 1, 2 * radius + 1)), iterations=1))
        
    # Determine the blur radius for each pixel based on the depth map
    blur_radii = np.clip(np.abs(disp - focus_disp) * max_blur_radius, 0, max_blur_radius).astype(np.int32)
    for radius in range(max_blur_radius + 1):
        dialted_blur_radii = np.clip(np.abs(dilated_disp[radius] - focus_disp) * max_blur_radius, 0, max_blur_radius).astype(np.int32)
        mask = (dialted_blur_radii >= radius) & (dialted_blur_radii >= blur_radii) & (dilated_disp[radius] > disp)
        blur_radii[mask] = dialted_blur_radii[mask]
    blur_radii = np.clip(blur_radii, 0, max_blur_radius)
    blur_radii = cv2.blur(blur_radii, (5, 5))

    # Precalculate the blured image for each blur radius
    unique_radii = np.unique(blur_radii)
    precomputed = {}
    for radius in range(max_blur_radius + 1):
        if radius not in unique_radii:
            continue
        precomputed[radius] = disk_blur(img, radius)
        
    # Composit the blured image for each pixel
    output = np.zeros_like(img)
    for r in unique_radii:
        mask = blur_radii == r
        output[mask] = precomputed[r][mask]
        
    return output