Add optional GPU-accelerated panorama depth merge

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

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)