Spaces:
Running on Zero
Running on Zero
| """Geometry-aware RGB-depth-camera transforms and blur primitives. | |
| Every geometric function updates normalized camera intrinsics together with RGB | |
| and depth. Mixed depth interpolation uses bilinear disparity on smooth surfaces | |
| and nearest sampling around geometric discontinuities to avoid flying points. | |
| The evaluation loader uses the geometry-preserving transforms in this module. | |
| """ | |
| from typing import Literal, Optional, Tuple | |
| import numpy as np | |
| import cv2 | |
| from PIL import Image | |
| import utils3d | |
| from scipy.signal import fftconvolve | |
| 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]: | |
| """Sample a valid target pinhole view inside a source camera frustum. | |
| Args: | |
| src_intrinsics: Normalized source camera matrix ``float [3,3]``. | |
| tgt_aspect: Target width divided by target height. | |
| center_augmentation: Fraction controlling random optical-axis movement. | |
| fov_range_absolute: Minimum/maximum target FoV in degrees. | |
| fov_range_relative: Multipliers limiting target FoV relative to source. | |
| rng: NumPy random generator used for FoV and center sampling. | |
| Returns: | |
| tgt_intrinsics: Normalized target camera matrix ``float32 [3,3]``. | |
| rotation: Camera-space rotation ``float32 [3,3]`` mapping source rays | |
| into the sampled target view. | |
| """ | |
| 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: Optional[np.ndarray] = None, | |
| transform: Optional[np.ndarray] = None, | |
| tgt_size: Optional[Tuple[int, int]] = None, | |
| interpolation: Literal['nearest', 'bilinear', 'lanczos'] = 'nearest', | |
| sparse_mask: Optional[np.ndarray] = None, | |
| ) -> np.ndarray: | |
| """Warp an image-like array through a normalized planar homography. | |
| Lanczos downsampling first reduces the source to avoid aliasing. Sparse | |
| nearest-neighbor input optionally uses mask-aware pre-resizing so isolated | |
| samples are not discarded. | |
| Args: | |
| src_map: Source array ``[H,W]`` or ``[H,W,C]``. | |
| transform: Normalized 3x3 homography satisfying | |
| ``p_target = transform @ p_source``. | |
| tgt_size: Output tuple ``(height,width)``. | |
| interpolation: ``nearest``, ``bilinear``, or ``lanczos``. | |
| sparse_mask: Optional boolean source support ``[H,W]`` for sparse | |
| nearest-neighbor maps. | |
| Returns: | |
| Warped array ``[H_t,W_t]`` or ``[H_t,W_t,C]``. | |
| """ | |
| 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 crop_resize_view( | |
| src_image: np.ndarray, | |
| src_depth: np.ndarray, | |
| src_intrinsics: np.ndarray, | |
| tgt_size: Tuple[int, int], | |
| rng: Optional[np.random.Generator] = None, | |
| random_crop: bool = True, | |
| image_interpolation: Literal['nearest', 'bilinear', 'lanczos'] = 'lanczos', | |
| depth_interpolation: Literal['nearest', 'mixed'] = 'nearest', | |
| ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: | |
| """Crop a source view and resize it to an exact target resolution. | |
| Large images are cropped directly at target size. Smaller images first take | |
| the largest crop matching target aspect ratio and then resize. Normalized | |
| intrinsics are updated for crop geometry; a full-image resize alone does not | |
| alter normalized values. | |
| Args: | |
| src_image: RGB uint8 image ``[H_s,W_s,3]``. | |
| src_depth: Depth map ``[H_s,W_s]`` with NaN/Inf semantics. | |
| src_intrinsics: Normalized source matrix ``[3,3]``. | |
| tgt_size: Output tuple ``(height,width)``. | |
| rng: Optional random generator for crop offsets. | |
| random_crop: Randomize offsets instead of center cropping. | |
| image_interpolation: RGB resampling method. | |
| depth_interpolation: ``nearest`` or edge-aware ``mixed``. | |
| Returns: | |
| image: RGB array ``[H_t,W_t,3]``. | |
| depth: Depth array ``float32 [H_t,W_t]``. | |
| intrinsics: Updated normalized matrix ``float32 [3,3]``. | |
| """ | |
| if depth_interpolation not in {'nearest', 'mixed'}: | |
| raise ValueError(f"crop_resize_view only supports nearest/mixed depth interpolation, got {depth_interpolation}.") | |
| tgt_height, tgt_width = tgt_size | |
| src_height, src_width = src_image.shape[:2] | |
| if src_height <= 0 or src_width <= 0 or tgt_height <= 0 or tgt_width <= 0: | |
| raise ValueError( | |
| f"Invalid source/target size: src=({src_height}, {src_width}), tgt=({tgt_height}, {tgt_width})." | |
| ) | |
| if src_width >= tgt_width and src_height >= tgt_height: | |
| crop_width = tgt_width | |
| crop_height = tgt_height | |
| else: | |
| tgt_aspect = tgt_width / tgt_height | |
| if src_width / src_height >= tgt_aspect: | |
| crop_height = src_height | |
| crop_width = max(1, min(src_width, int(np.floor(src_height * tgt_aspect)))) | |
| else: | |
| crop_width = src_width | |
| crop_height = max(1, min(src_height, int(np.floor(src_width / tgt_aspect)))) | |
| max_x = max(0, src_width - crop_width) | |
| max_y = max(0, src_height - crop_height) | |
| if random_crop: | |
| if rng is None: | |
| rng = np.random.default_rng() | |
| x0 = int(rng.integers(0, max_x + 1)) if max_x > 0 else 0 | |
| y0 = int(rng.integers(0, max_y + 1)) if max_y > 0 else 0 | |
| else: | |
| x0 = max_x // 2 | |
| y0 = max_y // 2 | |
| x1, y1 = x0 + crop_width, y0 + crop_height | |
| cropped_image = src_image[y0:y1, x0:x1] | |
| cropped_depth = src_depth[y0:y1, x0:x1] | |
| if crop_width != tgt_width or crop_height != tgt_height: | |
| if image_interpolation == 'lanczos': | |
| tgt_image = np.array( | |
| Image.fromarray(cropped_image).resize((tgt_width, tgt_height), Image.Resampling.LANCZOS) | |
| ) | |
| else: | |
| cv2_interpolation = { | |
| 'nearest': cv2.INTER_NEAREST, | |
| 'bilinear': cv2.INTER_LINEAR, | |
| }[image_interpolation] | |
| tgt_image = cv2.resize(cropped_image, (tgt_width, tgt_height), interpolation=cv2_interpolation) | |
| cropped_valid = np.isfinite(cropped_depth) | |
| cropped_depth_values = np.where(cropped_valid, cropped_depth, 0).astype(np.float32) | |
| tgt_depth_nearest = cv2.resize(cropped_depth_values, (tgt_width, tgt_height), interpolation=cv2.INTER_NEAREST) | |
| tgt_depth_valid = cv2.resize(cropped_valid.astype(np.uint8), (tgt_width, tgt_height), interpolation=cv2.INTER_NEAREST).astype(bool) | |
| if depth_interpolation == 'mixed': | |
| depth_edge_mask = utils3d.np.depth_map_edge(cropped_depth, mask=cropped_valid, kernel_size=5, ltol=0.005) | |
| depth_bilinear_mask = cropped_valid & ~depth_edge_mask | |
| tgt_depth_bilinear_mask = cv2.resize( | |
| depth_bilinear_mask.astype(np.float32), | |
| (tgt_width, tgt_height), | |
| interpolation=cv2.INTER_LINEAR, | |
| ) | |
| cropped_disp = np.where(cropped_valid & (cropped_depth > 0), 1.0 / cropped_depth, 0.0).astype(np.float32) | |
| tgt_disp_bilinear = cv2.resize(cropped_disp, (tgt_width, tgt_height), interpolation=cv2.INTER_LINEAR) | |
| tgt_depth_bilinear = np.where(tgt_disp_bilinear > 0, 1.0 / tgt_disp_bilinear, np.inf).astype(np.float32) | |
| tgt_depth = np.where(tgt_depth_bilinear_mask == 1.0, tgt_depth_bilinear, tgt_depth_nearest) | |
| else: | |
| tgt_depth = tgt_depth_nearest | |
| tgt_depth = np.where(tgt_depth_valid, tgt_depth, np.inf).astype(np.float32) | |
| else: | |
| tgt_image = cropped_image.copy() | |
| tgt_depth = cropped_depth.copy().astype(np.float32) | |
| tgt_intrinsics = src_intrinsics.astype(np.float32).copy() | |
| tgt_intrinsics[0, 0] = src_intrinsics[0, 0] * src_width / crop_width | |
| tgt_intrinsics[0, 1] = src_intrinsics[0, 1] * src_width / crop_width | |
| tgt_intrinsics[0, 2] = (src_intrinsics[0, 2] * src_width - x0) / crop_width | |
| tgt_intrinsics[1, 0] = 0.0 | |
| tgt_intrinsics[1, 1] = src_intrinsics[1, 1] * src_height / crop_height | |
| tgt_intrinsics[1, 2] = (src_intrinsics[1, 2] * src_height - y0) / crop_height | |
| tgt_intrinsics[2, 0] = 0.0 | |
| tgt_intrinsics[2, 1] = 0.0 | |
| tgt_intrinsics[2, 2] = 1.0 | |
| return tgt_image, tgt_depth, tgt_intrinsics | |
| def resize_then_crop_view( | |
| src_image: np.ndarray, | |
| src_depth: np.ndarray, | |
| src_intrinsics: np.ndarray, | |
| resize_size: Tuple[int, int], | |
| crop_size: Tuple[int, int], | |
| rng: Optional[np.random.Generator] = None, | |
| random_crop: bool = True, | |
| image_interpolation: Literal['nearest', 'bilinear', 'area'] = 'area', | |
| depth_interpolation: Literal['nearest'] = 'nearest', | |
| ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: | |
| """Resize the full view first, then crop an exact output window. | |
| Args: | |
| src_image: RGB uint8 image ``[H_s,W_s,3]``. | |
| src_depth: Depth map ``[H_s,W_s]``. | |
| src_intrinsics: Normalized source matrix ``[3,3]``. | |
| resize_size: Intermediate tuple ``(height,width)``. | |
| crop_size: Final tuple ``(height,width)`` no larger than resize size. | |
| rng: Optional random generator for crop offsets. | |
| random_crop: Randomize offsets instead of center cropping. | |
| image_interpolation: Intermediate RGB resampling method. | |
| depth_interpolation: Depth method; only nearest is supported. | |
| Returns: | |
| image: Cropped RGB ``[H_c,W_c,3]``. | |
| depth: Cropped depth ``float32 [H_c,W_c]``. | |
| intrinsics: Crop-adjusted normalized matrix ``float32 [3,3]``. | |
| """ | |
| if depth_interpolation != 'nearest': | |
| raise ValueError(f"resize_then_crop_view only supports nearest depth interpolation, got {depth_interpolation}.") | |
| resize_height, resize_width = resize_size | |
| crop_height, crop_width = crop_size | |
| src_height, src_width = src_image.shape[:2] | |
| if src_height <= 0 or src_width <= 0 or resize_height <= 0 or resize_width <= 0 or crop_height <= 0 or crop_width <= 0: | |
| raise ValueError( | |
| f"Invalid source/resize/crop size: src=({src_height}, {src_width}), " | |
| f"resize=({resize_height}, {resize_width}), crop=({crop_height}, {crop_width})." | |
| ) | |
| if crop_height > resize_height or crop_width > resize_width: | |
| raise ValueError( | |
| f"Crop size ({crop_height}, {crop_width}) must not exceed resize size ({resize_height}, {resize_width})." | |
| ) | |
| image_cv2_interpolation = { | |
| 'nearest': cv2.INTER_NEAREST, | |
| 'bilinear': cv2.INTER_LINEAR, | |
| 'area': cv2.INTER_AREA, | |
| }[image_interpolation] | |
| resized_image = cv2.resize(src_image, (resize_width, resize_height), interpolation=image_cv2_interpolation) | |
| resized_depth = cv2.resize(src_depth.astype(np.float32), (resize_width, resize_height), interpolation=cv2.INTER_NEAREST) | |
| max_x = max(0, resize_width - crop_width) | |
| max_y = max(0, resize_height - crop_height) | |
| if random_crop: | |
| if rng is None: | |
| rng = np.random.default_rng() | |
| x0 = int(rng.integers(0, max_x + 1)) if max_x > 0 else 0 | |
| y0 = int(rng.integers(0, max_y + 1)) if max_y > 0 else 0 | |
| else: | |
| x0 = max_x // 2 | |
| y0 = max_y // 2 | |
| x1, y1 = x0 + crop_width, y0 + crop_height | |
| tgt_image = resized_image[y0:y1, x0:x1].copy() | |
| tgt_depth = resized_depth[y0:y1, x0:x1].copy().astype(np.float32) | |
| tgt_intrinsics = src_intrinsics.astype(np.float32).copy() | |
| tgt_intrinsics[0, 0] = src_intrinsics[0, 0] * resize_width / crop_width | |
| tgt_intrinsics[0, 1] = src_intrinsics[0, 1] * resize_width / crop_width | |
| tgt_intrinsics[0, 2] = (src_intrinsics[0, 2] * resize_width - x0) / crop_width | |
| tgt_intrinsics[1, 0] = 0.0 | |
| tgt_intrinsics[1, 1] = src_intrinsics[1, 1] * resize_height / crop_height | |
| tgt_intrinsics[1, 2] = (src_intrinsics[1, 2] * resize_height - y0) / crop_height | |
| tgt_intrinsics[2, 0] = 0.0 | |
| tgt_intrinsics[2, 1] = 0.0 | |
| tgt_intrinsics[2, 2] = 1.0 | |
| return tgt_image, tgt_depth, tgt_intrinsics | |
| def resize_to_cover_center_crop_view( | |
| src_image: np.ndarray, | |
| src_depth: np.ndarray, | |
| src_intrinsics: np.ndarray, | |
| tgt_size: Tuple[int, int], | |
| image_interpolation: Literal['nearest', 'bilinear', 'lanczos'] = 'lanczos', | |
| depth_interpolation: Literal['nearest', 'mixed'] = 'mixed', | |
| ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: | |
| """Resize a view to cover the target, then take a centered crop. | |
| Args: | |
| src_image: RGB uint8 image ``[H_s,W_s,3]``. | |
| src_depth: Depth map ``[H_s,W_s]`` with invalid sentinels. | |
| src_intrinsics: Normalized source matrix ``[3,3]``. | |
| tgt_size: Output tuple ``(height,width)``. | |
| image_interpolation: RGB resampling method. | |
| depth_interpolation: ``nearest`` or edge-aware ``mixed``. | |
| Returns: | |
| image: Center-cropped RGB ``[H_t,W_t,3]``. | |
| depth: Center-cropped depth ``float32 [H_t,W_t]``. | |
| intrinsics: Resize/crop-adjusted normalized matrix ``float32 [3,3]``. | |
| """ | |
| if depth_interpolation not in {'nearest', 'mixed'}: | |
| raise ValueError( | |
| f"resize_to_cover_center_crop_view only supports nearest/mixed depth interpolation, got {depth_interpolation}." | |
| ) | |
| tgt_height, tgt_width = tgt_size | |
| src_height, src_width = src_image.shape[:2] | |
| if src_height <= 0 or src_width <= 0 or tgt_height <= 0 or tgt_width <= 0: | |
| raise ValueError( | |
| f"Invalid source/target size: src=({src_height}, {src_width}), tgt=({tgt_height}, {tgt_width})." | |
| ) | |
| scale = max(tgt_width / src_width, tgt_height / src_height) | |
| resized_width = max(tgt_width, int(round(src_width * scale))) | |
| resized_height = max(tgt_height, int(round(src_height * scale))) | |
| if image_interpolation == 'lanczos': | |
| resized_image = np.array( | |
| Image.fromarray(src_image).resize((resized_width, resized_height), Image.Resampling.LANCZOS) | |
| ) | |
| else: | |
| cv2_interpolation = { | |
| 'nearest': cv2.INTER_NEAREST, | |
| 'bilinear': cv2.INTER_LINEAR, | |
| }[image_interpolation] | |
| resized_image = cv2.resize(src_image, (resized_width, resized_height), interpolation=cv2_interpolation) | |
| resized_valid = np.isfinite(src_depth) | |
| resized_depth_values = cv2.resize( | |
| np.where(resized_valid, src_depth, 0).astype(np.float32), | |
| (resized_width, resized_height), | |
| interpolation=cv2.INTER_NEAREST, | |
| ) | |
| resized_depth_valid = cv2.resize( | |
| resized_valid.astype(np.uint8), | |
| (resized_width, resized_height), | |
| interpolation=cv2.INTER_NEAREST, | |
| ).astype(bool) | |
| resized_depth = np.where(resized_depth_valid, resized_depth_values, np.inf).astype(np.float32) | |
| if depth_interpolation == 'mixed': | |
| depth_edge_mask = utils3d.np.depth_map_edge(src_depth, mask=np.isfinite(src_depth), kernel_size=5, ltol=0.005) | |
| depth_bilinear_mask = np.isfinite(src_depth) & ~depth_edge_mask | |
| resized_bilinear_mask = cv2.resize( | |
| depth_bilinear_mask.astype(np.float32), | |
| (resized_width, resized_height), | |
| interpolation=cv2.INTER_LINEAR, | |
| ) | |
| disp = np.where(np.isfinite(src_depth) & (src_depth > 0), 1.0 / src_depth, 0.0).astype(np.float32) | |
| resized_disp = cv2.resize(disp, (resized_width, resized_height), interpolation=cv2.INTER_LINEAR) | |
| resized_depth_bilinear = np.where(resized_disp > 0, 1.0 / resized_disp, np.inf).astype(np.float32) | |
| resized_depth = np.where(resized_bilinear_mask == 1.0, resized_depth_bilinear, resized_depth) | |
| resized_depth = np.where(resized_depth_valid, resized_depth, np.inf).astype(np.float32) | |
| x0 = max(0, (resized_width - tgt_width) // 2) | |
| y0 = max(0, (resized_height - tgt_height) // 2) | |
| x1, y1 = x0 + tgt_width, y0 + tgt_height | |
| tgt_image = resized_image[y0:y1, x0:x1].copy() | |
| tgt_depth = resized_depth[y0:y1, x0:x1].copy().astype(np.float32) | |
| tgt_intrinsics = src_intrinsics.astype(np.float32).copy() | |
| tgt_intrinsics[0, 0] = src_intrinsics[0, 0] * resized_width / tgt_width | |
| tgt_intrinsics[0, 1] = src_intrinsics[0, 1] * resized_width / tgt_width | |
| tgt_intrinsics[0, 2] = (src_intrinsics[0, 2] * resized_width - x0) / tgt_width | |
| tgt_intrinsics[1, 0] = 0.0 | |
| tgt_intrinsics[1, 1] = src_intrinsics[1, 1] * resized_height / tgt_height | |
| tgt_intrinsics[1, 2] = (src_intrinsics[1, 2] * resized_height - y0) / tgt_height | |
| tgt_intrinsics[2, 0] = 0.0 | |
| tgt_intrinsics[2, 1] = 0.0 | |
| tgt_intrinsics[2, 2] = 1.0 | |
| return tgt_image, tgt_depth, tgt_intrinsics | |
| def resize_view( | |
| src_image: np.ndarray, | |
| src_depth: np.ndarray, | |
| src_intrinsics: np.ndarray, | |
| tgt_size: Tuple[int, int], | |
| image_interpolation: Literal['nearest', 'bilinear', 'lanczos'] = 'lanczos', | |
| depth_interpolation: Literal['nearest', 'mixed'] = 'mixed', | |
| ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: | |
| """Resize RGB and depth directly to a target width and height. | |
| Args: | |
| src_image: RGB uint8 image ``[H_s,W_s,3]``. | |
| src_depth: Depth map ``[H_s,W_s]`` with invalid sentinels. | |
| src_intrinsics: Normalized source matrix ``[3,3]``. | |
| tgt_size: Exact output tuple ``(height,width)``. | |
| image_interpolation: RGB resampling method. | |
| depth_interpolation: ``nearest`` or edge-aware ``mixed``. | |
| Returns: | |
| image: Resized RGB ``[H_t,W_t,3]``. | |
| depth: Resized depth ``float32 [H_t,W_t]``. | |
| intrinsics: Normalized matrix ``float32 [3,3]``. Direct full-image | |
| resizing leaves normalized focal lengths and center unchanged. | |
| """ | |
| if depth_interpolation not in {'nearest', 'mixed'}: | |
| raise ValueError(f"resize_view only supports nearest/mixed depth interpolation, got {depth_interpolation}.") | |
| tgt_height, tgt_width = tgt_size | |
| src_height, src_width = src_image.shape[:2] | |
| if src_height <= 0 or src_width <= 0 or tgt_height <= 0 or tgt_width <= 0: | |
| raise ValueError( | |
| f"Invalid source/target size: src=({src_height}, {src_width}), tgt=({tgt_height}, {tgt_width})." | |
| ) | |
| if image_interpolation == 'lanczos': | |
| tgt_image = np.array(Image.fromarray(src_image).resize((tgt_width, tgt_height), Image.Resampling.LANCZOS)) | |
| else: | |
| cv2_interpolation = { | |
| 'nearest': cv2.INTER_NEAREST, | |
| 'bilinear': cv2.INTER_LINEAR, | |
| }[image_interpolation] | |
| tgt_image = cv2.resize(src_image, (tgt_width, tgt_height), interpolation=cv2_interpolation) | |
| src_valid = np.isfinite(src_depth) | |
| tgt_depth_values = cv2.resize( | |
| np.where(src_valid, src_depth, 0).astype(np.float32), | |
| (tgt_width, tgt_height), | |
| interpolation=cv2.INTER_NEAREST, | |
| ) | |
| tgt_depth_valid = cv2.resize( | |
| src_valid.astype(np.uint8), | |
| (tgt_width, tgt_height), | |
| interpolation=cv2.INTER_NEAREST, | |
| ).astype(bool) | |
| tgt_depth = np.where(tgt_depth_valid, tgt_depth_values, np.inf).astype(np.float32) | |
| if depth_interpolation == 'mixed': | |
| depth_edge_mask = utils3d.np.depth_map_edge(src_depth, mask=src_valid, kernel_size=5, ltol=0.005) | |
| depth_bilinear_mask = src_valid & ~depth_edge_mask | |
| tgt_depth_bilinear_mask = cv2.resize( | |
| depth_bilinear_mask.astype(np.float32), | |
| (tgt_width, tgt_height), | |
| interpolation=cv2.INTER_LINEAR, | |
| ) | |
| src_disp = np.where(src_valid & (src_depth > 0), 1.0 / src_depth, 0.0).astype(np.float32) | |
| tgt_disp_bilinear = cv2.resize(src_disp, (tgt_width, tgt_height), interpolation=cv2.INTER_LINEAR) | |
| tgt_depth_bilinear = np.where(tgt_disp_bilinear > 0, 1.0 / tgt_disp_bilinear, np.inf).astype(np.float32) | |
| tgt_depth = np.where(tgt_depth_bilinear_mask == 1.0, tgt_depth_bilinear, tgt_depth) | |
| tgt_depth = np.where(tgt_depth_valid, tgt_depth, np.inf).astype(np.float32) | |
| tgt_intrinsics = src_intrinsics.astype(np.float32).copy() | |
| tgt_intrinsics[1, 0] = 0.0 | |
| tgt_intrinsics[2, 0] = 0.0 | |
| tgt_intrinsics[2, 1] = 0.0 | |
| tgt_intrinsics[2, 2] = 1.0 | |
| return tgt_image, tgt_depth, tgt_intrinsics | |
| def disk_kernel(radius: int) -> np.ndarray: | |
| """Generate a normalized circular convolution kernel. | |
| Args: | |
| radius: Nonnegative disk radius in pixels. | |
| Returns: | |
| Float32 kernel ``[2*radius+1,2*radius+1]`` summing to one. | |
| """ | |
| # 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 a circular point-spread function with FFT convolution. | |
| Args: | |
| image: Scalar ``[H,W]`` or channel image ``[H,W,C]``. | |
| radius: Nonnegative blur radius in pixels. | |
| Returns: | |
| Blurred floating array with the same shape as ``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: | |
| """Synthesize depth of field from a disparity map and focus plane. | |
| Args: | |
| img: RGB image ``[H,W,3]``. | |
| disp: Positive disparity map ``[H,W]`` aligned to ``img``. | |
| focus_disp: Disparity value lying on the simulated focus plane. | |
| max_blur_radius: Largest circular blur radius in pixels. | |
| Returns: | |
| Depth-of-field image with shape ``[H,W,3]`` and ``img`` dtype. | |
| """ | |
| # 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 | |