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 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