File size: 9,407 Bytes
ec0c8fa |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 |
from typing import *
from functools import partial
import math
import numpy as np
import utils3d
from .tools import timeit
def weighted_mean_numpy(x: np.ndarray, w: np.ndarray = None, axis: Union[int, Tuple[int,...]] = None, keepdims: bool = False, eps: float = 1e-7) -> np.ndarray:
if w is None:
return np.mean(x, axis=axis)
else:
w = w.astype(x.dtype)
return (x * w).mean(axis=axis) / np.clip(w.mean(axis=axis), eps, None)
def harmonic_mean_numpy(x: np.ndarray, w: np.ndarray = None, axis: Union[int, Tuple[int,...]] = None, keepdims: bool = False, eps: float = 1e-7) -> np.ndarray:
if w is None:
return 1 / (1 / np.clip(x, eps, None)).mean(axis=axis)
else:
w = w.astype(x.dtype)
return 1 / (weighted_mean_numpy(1 / (x + eps), w, axis=axis, keepdims=keepdims, eps=eps) + eps)
def image_plane_uv_numpy(width: int, height: int, aspect_ratio: float = None, dtype: np.dtype = np.float32) -> np.ndarray:
"UV with left-top corner as (-width / diagonal, -height / diagonal) and right-bottom corner as (width / diagonal, height / diagonal)"
if aspect_ratio is None:
aspect_ratio = width / height
span_x = aspect_ratio / (1 + aspect_ratio ** 2) ** 0.5
span_y = 1 / (1 + aspect_ratio ** 2) ** 0.5
u = np.linspace(-span_x * (width - 1) / width, span_x * (width - 1) / width, width, dtype=dtype)
v = np.linspace(-span_y * (height - 1) / height, span_y * (height - 1) / height, height, dtype=dtype)
u, v = np.meshgrid(u, v, indexing='xy')
uv = np.stack([u, v], axis=-1)
return uv
def focal_to_fov_numpy(focal: np.ndarray):
return 2 * np.arctan(0.5 / focal)
def fov_to_focal_numpy(fov: np.ndarray):
return 0.5 / np.tan(fov / 2)
def intrinsics_to_fov_numpy(intrinsics: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
fov_x = focal_to_fov_numpy(intrinsics[..., 0, 0])
fov_y = focal_to_fov_numpy(intrinsics[..., 1, 1])
return fov_x, fov_y
def solve_optimal_shift_focal(uv: np.ndarray, xyz: np.ndarray, ransac_iters: int = None, ransac_hypothetical_size: float = 0.1, ransac_threshold: float = 0.1):
"Solve `min |focal * xy / (z + shift) - uv|` with respect to shift and focal"
from scipy.optimize import least_squares
uv, xy, z = uv.reshape(-1, 2), xyz[..., :2].reshape(-1, 2), xyz[..., 2].reshape(-1)
def fn(uv: np.ndarray, xy: np.ndarray, z: np.ndarray, shift: np.ndarray):
xy_proj = xy / (z + shift)[: , None]
f = (xy_proj * uv).sum() / np.square(xy_proj).sum()
err = (f * xy_proj - uv).ravel()
return err
initial_shift = 0 #-z.min(keepdims=True) + 1.0
if ransac_iters is None:
solution = least_squares(partial(fn, uv, xy, z), x0=initial_shift, ftol=1e-3, method='lm')
optim_shift = solution['x'].squeeze().astype(np.float32)
else:
best_err, best_shift = np.inf, None
for _ in range(ransac_iters):
maybe_inliers = np.random.choice(len(z), size=int(ransac_hypothetical_size * len(z)), replace=False)
solution = least_squares(partial(fn, uv[maybe_inliers], xy[maybe_inliers], z[maybe_inliers]), x0=initial_shift, ftol=1e-3, method='lm')
maybe_shift = solution['x'].squeeze().astype(np.float32)
confirmed_inliers = np.linalg.norm(fn(uv, xy, z, maybe_shift).reshape(-1, 2), axis=-1) < ransac_threshold
if confirmed_inliers.sum() > 10:
solution = least_squares(partial(fn, uv[confirmed_inliers], xy[confirmed_inliers], z[confirmed_inliers]), x0=maybe_shift, ftol=1e-3, method='lm')
better_shift = solution['x'].squeeze().astype(np.float32)
else:
better_shift = maybe_shift
err = np.linalg.norm(fn(uv, xy, z, better_shift).reshape(-1, 2), axis=-1).clip(max=ransac_threshold).mean()
if err < best_err:
best_err, best_shift = err, better_shift
initial_shift = best_shift
optim_shift = best_shift
xy_proj = xy / (z + optim_shift)[: , None]
optim_focal = (xy_proj * uv).sum() / (xy_proj * xy_proj).sum()
return optim_shift, optim_focal
def point_map_to_depth_numpy(points: np.ndarray, mask: np.ndarray = None, downsample_size: Tuple[int, int] = (64, 64)):
import cv2
assert points.shape[-1] == 3, "Points should (H, W, 3)"
height, width = points.shape[-3], points.shape[-2]
diagonal = (height ** 2 + width ** 2) ** 0.5
uv = image_plane_uv_numpy(width=width, height=height)
if mask is None:
points_lr = cv2.resize(points, downsample_size, interpolation=cv2.INTER_LINEAR).reshape(-1, 3)
uv_lr = cv2.resize(uv, downsample_size, interpolation=cv2.INTER_LINEAR).reshape(-1, 2)
else:
index, mask_lr = mask_aware_nearest_resize_numpy(mask, *downsample_size)
points_lr, uv_lr = points[index][mask_lr], uv[index][mask_lr]
if points_lr.size == 0:
return np.zeros((height, width)), 0, 0, 0
optim_shift, optim_focal = solve_optimal_shift_focal(uv_lr, points_lr, ransac_iters=None)
fov_x = 2 * np.arctan(width / diagonal / optim_focal)
fov_y = 2 * np.arctan(height / diagonal / optim_focal)
depth = points[:, :, 2] + optim_shift
return depth, fov_x, fov_y, optim_shift
def mask_aware_nearest_resize_numpy(mask: np.ndarray, target_width: int, target_height: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""
Resize 2D map by nearest interpolation. Return the nearest neighbor index and mask of the resized map.
### Parameters
- `mask`: Input 2D mask of shape (..., H, W)
- `target_width`: target width of the resized map
- `target_height`: target height of the resized map
### Returns
- `nearest_idx`: Nearest neighbor index of the resized map of shape (..., target_height, target_width). Indices are like j + i * W, where j is the row index and i is the column index.
- `target_mask`: Mask of the resized map of shape (..., target_height, target_width)
"""
height, width = mask.shape[-2:]
filter_h_f, filter_w_f = max(1, height / target_height), max(1, width / target_width)
filter_h_i, filter_w_i = math.ceil(filter_h_f), math.ceil(filter_w_f)
filter_size = filter_h_i * filter_w_i
padding_h, padding_w = round(filter_h_f / 2), round(filter_w_f / 2)
# Window the original mask and uv
uv = utils3d.numpy.image_pixel_center(width=width, height=height, dtype=np.float32)
indices = np.arange(height * width, dtype=np.int32).reshape(height, width)
padded_uv = np.full((height + 2 * padding_h, width + 2 * padding_w, 2), 0, dtype=np.float32)
padded_uv[padding_h:padding_h + height, padding_w:padding_w + width] = uv
padded_mask = np.full((*mask.shape[:-2], height + 2 * padding_h, width + 2 * padding_w), False, dtype=bool)
padded_mask[..., padding_h:padding_h + height, padding_w:padding_w + width] = mask
padded_indices = np.full((height + 2 * padding_h, width + 2 * padding_w), 0, dtype=np.int32)
padded_indices[padding_h:padding_h + height, padding_w:padding_w + width] = indices
windowed_uv = utils3d.numpy.sliding_window_2d(padded_uv, (filter_h_i, filter_w_i), 1, axis=(0, 1))
windowed_mask = utils3d.numpy.sliding_window_2d(padded_mask, (filter_h_i, filter_w_i), 1, axis=(-2, -1))
windowed_indices = utils3d.numpy.sliding_window_2d(padded_indices, (filter_h_i, filter_w_i), 1, axis=(0, 1))
# Gather the target pixels's local window
target_uv = utils3d.numpy.image_uv(width=target_width, height=target_height, dtype=np.float32) * np.array([width, height], dtype=np.float32)
target_corner = target_uv - np.array((filter_w_f / 2, filter_h_f / 2), dtype=np.float32)
target_corner = np.round(target_corner - 0.5).astype(np.int32) + np.array((padding_w, padding_h), dtype=np.int32)
target_window_uv = windowed_uv[target_corner[..., 1], target_corner[..., 0], :, :, :].reshape(target_height, target_width, 2, filter_size) # (target_height, tgt_width, 2, filter_size)
target_window_mask = windowed_mask[..., target_corner[..., 1], target_corner[..., 0], :, :].reshape(*mask.shape[:-2], target_height, target_width, filter_size) # (..., target_height, tgt_width, filter_size)
target_window_indices = windowed_indices[target_corner[..., 1], target_corner[..., 0], :, :].reshape(target_height, target_width, filter_size) # (target_height, tgt_width, filter_size)
# Compute nearest neighbor in the local window for each pixel
dist = np.square(target_window_uv - target_uv[..., None])
dist = dist[..., 0, :] + dist[..., 1, :]
dist = np.where(target_window_mask, dist, np.inf) # (..., target_height, tgt_width, filter_size)
nearest_in_window = np.argmin(dist, axis=-1, keepdims=True) # (..., target_height, tgt_width, 1)
nearest_idx = np.take_along_axis(target_window_indices, nearest_in_window, axis=-1).squeeze(-1) # (..., target_height, tgt_width)
nearest_i, nearest_j = nearest_idx // width, nearest_idx % width
target_mask = np.any(target_window_mask, axis=-1)
batch_indices = [np.arange(n).reshape([1] * i + [n] + [1] * (mask.ndim - i - 1)) for i, n in enumerate(mask.shape[:-2])]
return (*batch_indices, nearest_i, nearest_j), target_mask |