diff --git a/moge/scripts/infer_panorama.py b/moge/scripts/infer_panorama.py index 0e93878..0c887f6 100644 --- a/moge/scripts/infer_panorama.py +++ b/moge/scripts/infer_panorama.py @@ -21,6 +21,7 @@ @click.option('--resolution_level', type=int, default=9, help='An integer [0-9] for the resolution level of inference. The higher, the better but slower. Defaults to 9. Note that it is irrelevant to the output resolution.') @click.option('--threshold', type=float, default=0.03, help='Threshold for removing edges. Defaults to 0.03. Smaller value removes more edges. "inf" means no thresholding.') @click.option('--batch_size', type=int, default=4, help='Batch size for inference. Defaults to 4.') +@click.option('--gpu_merge', is_flag=True, help='Solve the panorama depth merge on GPU instead of CPU.') @click.option('--splitted', 'save_splitted', is_flag=True, help='Whether to save the splitted images. Defaults to False.') @click.option('--maps', 'save_maps_', is_flag=True, help='Whether to save the output maps and fov(image, depth, mask, points, fov).') @click.option('--glb', 'save_glb_', is_flag=True, help='Whether to save the output as a.glb file. The color will be saved as a texture.') @@ -35,6 +36,7 @@ def main( resolution_level: int, threshold: float, batch_size: int, + gpu_merge: bool, save_splitted: bool, save_maps_: bool, save_glb_: bool, @@ -118,7 +120,11 @@ def main( print('Merging...') if pbar.disable else pbar.set_postfix_str(f'Merging') merging_width, merging_height = min(1920, width), min(960, height) - panorama_depth, panorama_mask = merge_panorama_depth(merging_width, merging_height, splitted_distance_maps, splitted_masks, splitted_extrinsics, splitted_intriniscs) + panorama_depth, panorama_mask = merge_panorama_depth( + merging_width, merging_height, splitted_distance_maps, splitted_masks, + splitted_extrinsics, splitted_intriniscs, + device=device_name if gpu_merge else None, + ) panorama_depth = panorama_depth.astype(np.float32) panorama_depth = cv2.resize(panorama_depth, (width, height), cv2.INTER_LINEAR) panorama_mask = cv2.resize(panorama_mask.astype(np.uint8), (width, height), cv2.INTER_NEAREST) > 0 diff --git a/moge/utils/panorama.py b/moge/utils/panorama.py index c593748..a21733c 100644 --- a/moge/utils/panorama.py +++ b/moge/utils/panorama.py @@ -105,9 +105,81 @@ def grad_equation(width: int, height: int, wrap_x: bool = False, wrap_y: bool = return A -def merge_panorama_depth(width: int, height: int, distance_maps: List[np.ndarray], pred_masks: List[np.ndarray], extrinsics: List[np.ndarray], intrinsics: List[np.ndarray]): +def _solve_gpu(A, b, device, iters: int = 100, tol: float = 1e-5, x0: Optional[np.ndarray] = None) -> np.ndarray: + import torch + + if isinstance(device, str): + device = torch.device(device) + + A_csr = A.tocsr() + A_csr.sort_indices() + AT_csr = A.T.tocsr() + AT_csr.sort_indices() + + idx_dtype = torch.int32 if A_csr.nnz < (2**31 - 1) else torch.int64 + + A_t = torch.sparse_csr_tensor( + torch.as_tensor(A_csr.indptr, dtype=idx_dtype, device=device), + torch.as_tensor(A_csr.indices, dtype=idx_dtype, device=device), + torch.as_tensor(A_csr.data, dtype=torch.float32, device=device), + size=A_csr.shape, + device=device, + ) + AT_t = torch.sparse_csr_tensor( + torch.as_tensor(AT_csr.indptr, dtype=idx_dtype, device=device), + torch.as_tensor(AT_csr.indices, dtype=idx_dtype, device=device), + torch.as_tensor(AT_csr.data, dtype=torch.float32, device=device), + size=AT_csr.shape, + device=device, + ) + b_t = torch.as_tensor(b, dtype=torch.float32, device=device) + + def spmm(mat, vec): + return torch.sparse.mm(mat, vec.unsqueeze(1)).squeeze(1) + + if x0 is not None: + x = torch.as_tensor(x0, dtype=torch.float32, device=device) + e = b_t - spmm(A_t, x) + else: + x = torch.zeros(A.shape[1], dtype=torch.float32, device=device) + e = b_t + + r = spmm(AT_t, e) + p = r.clone() + rs_old = torch.dot(r, r) + + if torch.sqrt(rs_old) < tol: + return x.cpu().numpy() + + for _ in range(iters): + q = spmm(A_t, p) + pq = torch.dot(q, q) + if pq < 1e-12: + break + alpha = rs_old / pq + x += alpha * p + e -= alpha * q + r = spmm(AT_t, e) + rs_new = torch.dot(r, r) + if torch.sqrt(rs_new) < tol: + break + p = r + (rs_new / rs_old) * p + rs_old = rs_new + + return x.cpu().numpy() + + +def merge_panorama_depth( + width: int, + height: int, + distance_maps: List[np.ndarray], + pred_masks: List[np.ndarray], + extrinsics: List[np.ndarray], + intrinsics: List[np.ndarray], + device: Optional[Union[str, Any]] = None, +): if max(width, height) > 256: - panorama_depth_init, _ = merge_panorama_depth(width // 2, height // 2, distance_maps, pred_masks, extrinsics, intrinsics) + panorama_depth_init, _ = merge_panorama_depth(width // 2, height // 2, distance_maps, pred_masks, extrinsics, intrinsics, device=device) panorama_depth_init = cv2.resize(panorama_depth_init, (width, height), cv2.INTER_LINEAR) else: panorama_depth_init = None @@ -180,12 +252,16 @@ def merge_panorama_depth(width: int, height: int, distance_maps: List[np.ndarray panorama_log_distance_grad_y.reshape(-1)[grad_y_mask], panorama_laplacian_map.reshape(-1)[laplacian_mask] ]) - x, *_ = lsmr( - A, b, - atol=1e-5, btol=1e-5, - x0=np.log(panorama_depth_init).reshape(-1) if panorama_depth_init is not None else None, - show=False, - ) + x0 = np.log(panorama_depth_init).reshape(-1) if panorama_depth_init is not None else None + if device is None: + x, *_ = lsmr( + A, b, + atol=1e-5, btol=1e-5, + x0=x0, + show=False, + ) + else: + x = _solve_gpu(A, b, device=device, x0=x0) panorama_depth = np.exp(x).reshape(height, width).astype(np.float32) panorama_mask = np.any(panorama_pred_masks, axis=0)