origin

digit_depth/third_party/__init__.py

@@ -0,0 +1,4 @@
+from .data_utils import *
+from .geom_utils import *
+from .vis_utils import *
+from .poisson import *

digit_depth/third_party/data_utils.py

@@ -0,0 +1,30 @@
+import numpy as np
+import torch
+
+"""
+dataloader, logger helpers
+"""
+
+
+def pandas_col_to_numpy(df_col):
+
+    df_col = df_col.apply(lambda x: np.fromstring(x.replace("\n", "").replace("[", "").replace("]", "").replace("  ", " "), sep=", "))
+    df_col = np.stack(df_col)
+    return df_col
+
+
+def pandas_string_to_numpy(arr_str):
+    arr_npy = np.fromstring(arr_str.replace("\n", "").replace("[", "").replace("]", "").replace("  ", " "),sep=", ")
+    return arr_npy
+
+
+def interpolate_img(img, rows, cols):
+    """
+    img: C x H x W
+    """
+
+    img = torch.nn.functional.interpolate(img, size=cols)
+    img = torch.nn.functional.interpolate(img.permute(0, 2, 1), size=rows)
+    img = img.permute(0, 2, 1)
+
+    return img

digit_depth/third_party/geom_utils.py

@@ -0,0 +1,531 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+import copy
+
+import numpy as np
+import open3d as o3d
+import torch
+from scipy import ndimage
+
+from digit_depth.third_party.poisson import poisson_reconstruct
+from digit_depth.third_party.vis_utils import visualize_inlier_outlier
+
+"""
+Common functions
+"""
+
+
+def flip(x):
+    return torch.flip(x, dims=[0])
+
+
+def min_clip(x, min_val):
+    return torch.max(x, min_val)
+
+
+def max_clip(x, max_val):
+    return torch.min(x, max_val)
+
+
+def normalize(x, min_val, max_val):
+    return (x - torch.min(x)) * (max_val - min_val) / (
+        torch.max(x) - torch.min(x)
+    ) + min_val
+
+
+def mask_background(x, bg_mask, bg_val=0.0):
+    if bg_mask is not None:
+        x[bg_mask] = bg_val
+
+    return x
+
+
+def remove_background_pts(pts, bg_mask=None):
+    if bg_mask is not None:
+        fg_mask_pts = ~bg_mask.view(-1)
+        points3d_x = pts[0, fg_mask_pts].view(1, -1)
+        points3d_y = pts[1, fg_mask_pts].view(1, -1)
+        points3d_z = pts[2, fg_mask_pts].view(1, -1)
+        pts_fg = torch.cat((points3d_x, points3d_y, points3d_z), dim=0)
+    else:
+        pts_fg = pts
+
+    return pts_fg
+
+
+"""
+3D transform helper functions
+"""
+
+
+def Rt_to_T(R=None, t=None, device=None):
+    """
+    :param R: rotation, (B, 3, 3) or (3, 3)
+    :param t: translation, (B, 3) or  (3)
+    :return: T, (B, 4, 4) or  (4, 4)
+    """
+
+    T = torch.eye(4, device=device)
+
+    if (len(R.shape) > 2) & (len(t.shape) > 1):  # batch version
+        B = R.shape[0]
+        T = T.repeat(B, 1, 1)
+        if R is not None:
+            T[:, 0:3, 0:3] = R
+        if t is not None:
+            T[:, 0:3, -1] = t
+
+    else:
+        if R is not None:
+            T[0:3, 0:3] = R
+        if t is not None:
+            T[0:3, -1] = t
+
+    return T
+
+
+def transform_pts3d(T, pts):
+    """
+    T: 4x4
+    pts: 3xN
+
+    returns 3xN
+    """
+    D, N = pts.shape
+
+    if D == 3:
+        pts_tf = (
+            torch.cat((pts, torch.ones(1, N)), dim=0)
+            if torch.is_tensor(pts)
+            else np.concatenate((pts, torch.ones(1, N)), axis=0)
+        )
+
+    pts_tf = torch.matmul(T, pts_tf) if torch.is_tensor(pts) else np.matmul(T, pts_tf)
+    pts_tf = pts_tf[0:3, :]
+
+    return pts_tf
+
+
+"""
+3D-2D projection / conversion functions
+OpenGL transforms reference: http://www.songho.ca/opengl/gl_transform.html
+"""
+
+
+def _vectorize_pixel_coords(rows, cols, device=None):
+    y_range = torch.arange(rows, device=device)
+    x_range = torch.arange(cols, device=device)
+    grid_x, grid_y = torch.meshgrid(x_range, y_range)
+    pixel_pos = torch.stack([grid_x.reshape(-1), grid_y.reshape(-1)], dim=0)  # 2 x N
+
+    return pixel_pos
+
+
+def _clip_to_pixel(clip_pos, img_shape, params):
+    H, W = img_shape
+
+    # clip -> ndc coords
+    x_ndc = clip_pos[0, :] / clip_pos[3, :]
+    y_ndc = clip_pos[1, :] / clip_pos[3, :]
+    # z_ndc = clip_pos[2, :] / clip_pos[3, :]
+
+    # ndc -> pixel coords
+    x_pix = (W - 1) / 2 * (x_ndc + 1)  # [-1, 1] -> [0, W-1]
+    y_pix = (H - 1) / 2 * (y_ndc + 1)  # [-1, 1] -> [0, H-1]
+    # z_pix = (f-n) / 2 *  z_ndc + (f+n) / 2
+
+    pixel_pos = torch.stack((x_pix, y_pix), dim=0)
+
+    return pixel_pos
+
+
+def _pixel_to_clip(pix_pos, depth_map, params):
+    """
+    :param pix_pos: position in pixel space, (2, N)
+    :param depth_map: depth map, (H, W)
+    :return: clip_pos position in clip space, (4, N)
+    """
+
+    x_pix = pix_pos[0, :]
+    y_pix = pix_pos[1, :]
+
+    H, W = depth_map.shape
+    f = params.z_far
+    n = params.z_near
+
+    # pixel -> ndc coords
+    x_ndc = 2 / (W - 1) * x_pix - 1  # [0, W-1] -> [-1, 1]
+    y_ndc = 2 / (H - 1) * y_pix - 1  # [0, H-1] -> [-1, 1]
+    z_buf = depth_map[y_pix, x_pix]
+
+    # ndc -> clip coords
+    z_eye = -z_buf
+    w_c = -z_eye
+    x_c = x_ndc * w_c
+    y_c = y_ndc * w_c
+    z_c = -(f + n) / (f - n) * z_eye - 2 * f * n / (f - n) * 1.0
+
+    clip_pos = torch.stack([x_c, y_c, z_c, w_c], dim=0)
+
+    return clip_pos
+
+
+def _clip_to_eye(clip_pos, P):
+    P_inv = torch.inverse(P)
+    eye_pos = torch.matmul(P_inv, clip_pos)
+
+    return eye_pos
+
+
+def _eye_to_clip(eye_pos, P):
+    clip_pos = torch.matmul(P, eye_pos)
+
+    return clip_pos
+
+
+def _eye_to_world(eye_pos, V):
+    V_inv = torch.inverse(V)
+    world_pos = torch.matmul(V_inv, eye_pos)
+
+    world_pos = world_pos / world_pos[3, :]
+
+    return world_pos
+
+
+def _world_to_eye(world_pos, V):
+    eye_pos = torch.matmul(V, world_pos)
+
+    return eye_pos
+
+
+def _world_to_object(world_pos, M):
+    M_inv = torch.inverse(M)
+    obj_pos = torch.matmul(M_inv, world_pos)
+
+    obj_pos = obj_pos / obj_pos[3, :]
+
+    return obj_pos
+
+
+def _object_to_world(obj_pos, M):
+    world_pos = torch.matmul(M, obj_pos)
+
+    world_pos = world_pos / world_pos[3, :]
+
+    return world_pos
+
+
+def depth_to_pts3d(depth, P, V, params=None, ordered_pts=False):
+    """
+    :param depth: depth map, (C, H, W) or (H, W)
+    :param P: projection matrix, (4, 4)
+    :param V: view matrix, (4, 4)
+    :return: world_pos position in 3d world coordinates, (3, H, W) or (3, N)
+    """
+
+    assert 2 <= len(depth.shape) <= 3
+    assert P.shape == (4, 4)
+    assert V.shape == (4, 4)
+
+    depth_map = depth.squeeze(0) if (len(depth.shape) == 3) else depth
+    H, W = depth_map.shape
+    pixel_pos = _vectorize_pixel_coords(rows=H, cols=W)
+
+    clip_pos = _pixel_to_clip(pixel_pos, depth_map, params)
+    eye_pos = _clip_to_eye(clip_pos, P)
+    world_pos = _eye_to_world(eye_pos, V)
+
+    world_pos = world_pos[0:3, :] / world_pos[3, :]
+
+    if ordered_pts:
+        H, W = depth_map.shape
+        world_pos = world_pos.reshape(world_pos.shape[0], H, W)
+
+    return world_pos
+
+
+"""
+Optical flow functions
+"""
+
+
+def analytic_flow(img1, depth1, P, V1, V2, M1, M2, gel_depth, params):
+    C, H, W = img1.shape
+    depth_map = depth1.squeeze(0) if (len(depth1.shape) == 3) else depth1
+    pixel_pos = _vectorize_pixel_coords(rows=H, cols=W, device=img1.device)
+
+    clip_pos = _pixel_to_clip(pixel_pos, depth_map, params)
+    eye_pos = _clip_to_eye(clip_pos, P)
+    world_pos = _eye_to_world(eye_pos, V1)
+    obj_pos = _world_to_object(world_pos, M1)
+
+    world_pos = _object_to_world(obj_pos, M2)
+    eye_pos = _world_to_eye(world_pos, V2)
+    clip_pos = _eye_to_clip(eye_pos, P)
+    pixel_pos_proj = _clip_to_pixel(clip_pos, (H, W), params)
+
+    pixel_flow = pixel_pos - pixel_pos_proj
+    flow_map = pixel_flow.reshape(pixel_flow.shape[0], H, W)
+
+    # mask out background gel pixels
+    mask_idxs = depth_map >= gel_depth
+    flow_map[:, mask_idxs] = 0.0
+
+    return flow_map
+
+
+"""
+Normal to depth conversion / integration functions
+"""
+
+
+def _preproc_depth(img_depth, bg_mask=None):
+    if bg_mask is not None:
+        img_depth = mask_background(img_depth, bg_mask=bg_mask, bg_val=0.0)
+
+    return img_depth
+
+
+def _preproc_normal(img_normal, bg_mask=None):
+    """
+    img_normal: lies in range [0, 1]
+    """
+
+    # 0.5 corresponds to 0
+    img_normal = img_normal - 0.5
+
+    # normalize
+    img_normal = img_normal / torch.linalg.norm(img_normal, dim=0)
+
+    # set background to have only z normals (flat, facing camera)
+    if bg_mask is not None:
+        img_normal[0:2, bg_mask] = 0.0
+        img_normal[2, bg_mask] = 1.0
+
+    return img_normal
+
+
+def _depth_to_grad_depth(img_depth, bg_mask=None):
+    gradx = ndimage.sobel(img_depth.cpu().detach().numpy(), axis=1, mode="constant")
+    grady = ndimage.sobel(img_depth.cpu().detach().numpy(), axis=0, mode="constant")
+
+    gradx = torch.FloatTensor(gradx, device=img_depth.device)
+    grady = torch.FloatTensor(grady, device=img_depth.device)
+
+    if bg_mask is not None:
+        gradx = mask_background(gradx, bg_mask=bg_mask, bg_val=0.0)
+        grady = mask_background(grady, bg_mask=bg_mask, bg_val=0.0)
+
+    return gradx, grady
+
+
+def _normal_to_grad_depth(img_normal, gel_width=1.0, gel_height=1.0, bg_mask=None):
+    # Ref: https://stackoverflow.com/questions/34644101/calculate-surface-normals-from-depth-image-using-neighboring-pixels-cross-produc/34644939#34644939
+
+    EPS = 1e-1
+    nz = torch.max(torch.tensor(EPS), img_normal[2, :])
+
+    dzdx = -(img_normal[0, :] / nz).squeeze()
+    dzdy = -(img_normal[1, :] / nz).squeeze()
+
+    # taking out negative sign as we are computing gradient of depth not z
+    # since z is pointed towards sensor, increase in z corresponds to decrease in depth
+    # i.e., dz/dx = -ddepth/dx
+    ddepthdx = -dzdx
+    ddepthdy = -dzdy
+
+    # sim: pixel axis v points in opposite dxn of camera axis y
+    ddepthdu = ddepthdx
+    ddepthdv = -ddepthdy
+
+    gradx = ddepthdu  # cols
+    grady = ddepthdv  # rows
+
+    # convert units from pixel to meters
+    C, H, W = img_normal.shape
+    gradx = gradx * (gel_width / W)
+    grady = grady * (gel_height / H)
+
+    if bg_mask is not None:
+        gradx = mask_background(gradx, bg_mask=bg_mask, bg_val=0.0)
+        grady = mask_background(grady, bg_mask=bg_mask, bg_val=0.0)
+
+    return gradx, grady
+
+
+def _integrate_grad_depth(gradx, grady, boundary=None, bg_mask=None, max_depth=0.0):
+    if boundary is None:
+        boundary = torch.zeros((gradx.shape[0], gradx.shape[1]))
+
+    img_depth_recon = poisson_reconstruct(
+        grady.cpu().detach().numpy(),
+        gradx.cpu().detach().numpy(),
+        boundary.cpu().detach().numpy(),
+    )
+    img_depth_recon = torch.tensor(
+        img_depth_recon, device=gradx.device, dtype=torch.float32
+    )
+
+    if bg_mask is not None:
+        img_depth_recon = mask_background(img_depth_recon, bg_mask)
+
+    # after integration, img_depth_recon lies between 0. (bdry) and a -ve val (obj depth)
+    # rescale to make max depth as gel depth and obj depth as +ve values
+    img_depth_recon = max_clip(img_depth_recon, max_val=torch.tensor(0.0)) + max_depth
+
+    return img_depth_recon
+
+
+def depth_to_depth(img_depth, bg_mask=None, boundary=None, params=None):
+    # preproc depth img
+    img_depth = _preproc_depth(img_depth=img_depth.squeeze(), bg_mask=bg_mask)
+
+    # get grad depth
+    gradx, grady = _depth_to_grad_depth(img_depth=img_depth.squeeze(), bg_mask=bg_mask)
+
+    # integrate grad depth
+    img_depth_recon = _integrate_grad_depth(
+        gradx, grady, boundary=boundary, bg_mask=bg_mask
+    )
+
+    return img_depth_recon
+
+
+def normal_to_depth(
+    img_normal,
+    bg_mask=None,
+    boundary=None,
+    gel_width=0.02,
+    gel_height=0.03,
+    max_depth=0.02,
+):
+    # preproc normal img
+    img_normal = _preproc_normal(img_normal=img_normal, bg_mask=bg_mask)
+
+    # get grad depth
+    gradx, grady = _normal_to_grad_depth(
+        img_normal=img_normal,
+        gel_width=gel_width,
+        gel_height=gel_height,
+        bg_mask=bg_mask,
+    )
+
+    # integrate grad depth
+    img_depth_recon = _integrate_grad_depth(
+        gradx, grady, boundary=boundary, bg_mask=bg_mask, max_depth=max_depth
+    )
+
+    return img_depth_recon
+
+
+"""
+3D registration functions
+"""
+
+
+def _fpfh(pcd, normals):
+    pcd.normals = o3d.utility.Vector3dVector(normals)
+    pcd_fpfh = o3d.pipelines.registration.compute_fpfh_feature(
+        pcd, o3d.geometry.KDTreeSearchParamHybrid(radius=0.3, max_nn=64)
+    )
+    return pcd_fpfh
+
+
+def _fast_global_registration(source, target, source_fpfh, target_fpfh):
+    distance_threshold = 0.01
+    result = o3d.pipelines.registration.registration_fast_based_on_feature_matching(
+        source,
+        target,
+        source_fpfh,
+        target_fpfh,
+        o3d.pipelines.registration.FastGlobalRegistrationOption(
+            maximum_correspondence_distance=distance_threshold
+        ),
+    )
+
+    transformation = result.transformation
+    metrics = [result.fitness, result.inlier_rmse, result.correspondence_set]
+
+    return transformation, metrics
+
+
+def fgr(source, target, src_normals, tgt_normals):
+    source_fpfh = _fpfh(source, src_normals)
+    target_fpfh = _fpfh(target, tgt_normals)
+    transformation, metrics = _fast_global_registration(
+        source=source, target=target, source_fpfh=source_fpfh, target_fpfh=target_fpfh
+    )
+    return transformation, metrics
+
+
+def icp(source, target, T_init=np.eye(4), mcd=0.1, max_iter=50, type="point_to_plane"):
+    if type == "point_to_point":
+        result = o3d.pipelines.registration.registration_icp(
+            source=source,
+            target=target,
+            max_correspondence_distance=mcd,
+            init=T_init,
+            estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPoint(),
+            criteria=o3d.pipelines.registration.ICPConvergenceCriteria(
+                max_iteration=max_iter
+            ),
+        )
+    else:
+        result = o3d.pipelines.registration.registration_icp(
+            source=source,
+            target=target,
+            max_correspondence_distance=mcd,
+            init=T_init,
+            estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPlane(),
+            criteria=o3d.pipelines.registration.ICPConvergenceCriteria(
+                max_iteration=max_iter
+            ),
+        )
+
+    transformation = result.transformation
+    metrics = [result.fitness, result.inlier_rmse, result.correspondence_set]
+
+    return transformation, metrics
+
+
+"""
+Open3D helper functions
+"""
+
+
+def remove_outlier_pts(points3d, nb_neighbors=20, std_ratio=10.0, vis=False):
+    points3d_np = (
+        points3d.cpu().detach().numpy() if torch.is_tensor(points3d) else points3d
+    )
+
+    cloud = o3d.geometry.PointCloud()
+    cloud.points = copy.deepcopy(o3d.utility.Vector3dVector(points3d_np.transpose()))
+    cloud_filt, ind_filt = cloud.remove_statistical_outlier(
+        nb_neighbors=nb_neighbors, std_ratio=std_ratio
+    )
+
+    if vis:
+        visualize_inlier_outlier(cloud, ind_filt)
+
+    points3d_filt = np.asarray(cloud_filt.points).transpose()
+    points3d_filt = (
+        torch.tensor(points3d_filt) if torch.is_tensor(points3d) else points3d_filt
+    )
+
+    return points3d_filt
+
+
+def init_points_to_clouds(clouds, points3d, colors=None):
+    for idx, cloud in enumerate(clouds):
+        points3d_np = (
+            points3d[idx].cpu().detach().numpy()
+            if torch.is_tensor(points3d[idx])
+            else points3d[idx]
+        )
+        cloud.points = copy.deepcopy(
+            o3d.utility.Vector3dVector(points3d_np.transpose())
+        )
+        if colors is not None:
+            cloud.paint_uniform_color(colors[idx])
+
+    return clouds

digit_depth/third_party/poisson.py

@@ -0,0 +1,81 @@
+"""
+poisson_reconstruct.py
+Fast Poisson Reconstruction in Python
+
+Copyright (c) 2014 Jack Doerner
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+"""
+
+import copy
+import math
+
+import numpy
+import scipy
+import scipy.fftpack
+
+
+def poisson_reconstruct(grady, gradx, boundarysrc):
+    # Thanks to Dr. Ramesh Raskar for providing the original matlab code from which this is derived
+    # Dr. Raskar's version is available here: http://web.media.mit.edu/~raskar/photo/code.pdf
+
+    # Laplacian
+    gyy = grady[1:, :-1] - grady[:-1, :-1]
+    gxx = gradx[:-1, 1:] - gradx[:-1, :-1]
+    f = numpy.zeros(boundarysrc.shape)
+    f[:-1, 1:] += gxx
+    f[1:, :-1] += gyy
+
+    # Boundary image
+    boundary = copy.deepcopy(boundarysrc)  # .copy()
+    boundary[1:-1, 1:-1] = 0
+
+    # Subtract boundary contribution
+    f_bp = (
+        -4 * boundary[1:-1, 1:-1]
+        + boundary[1:-1, 2:]
+        + boundary[1:-1, 0:-2]
+        + boundary[2:, 1:-1]
+        + boundary[0:-2, 1:-1]
+    )
+    f = f[1:-1, 1:-1] - f_bp
+
+    # Discrete Sine Transform
+    tt = scipy.fftpack.dst(f, norm="ortho")
+    fsin = scipy.fftpack.dst(tt.T, norm="ortho").T
+
+    # Eigenvalues
+    (x, y) = numpy.meshgrid(
+        range(1, f.shape[1] + 1), range(1, f.shape[0] + 1), copy=True
+    )
+    denom = (2 * numpy.cos(math.pi * x / (f.shape[1] + 2)) - 2) + (
+        2 * numpy.cos(math.pi * y / (f.shape[0] + 2)) - 2
+    )
+
+    f = fsin / denom
+
+    # Inverse Discrete Sine Transform
+    tt = scipy.fftpack.idst(f, norm="ortho")
+    img_tt = scipy.fftpack.idst(tt.T, norm="ortho").T
+
+    # New center + old boundary
+    result = copy.deepcopy(boundary)
+    result[1:-1, 1:-1] = img_tt
+
+    return result

digit_depth/third_party/vis_utils.py

@@ -0,0 +1,372 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+import copy
+import logging
+import time
+import types
+
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+import open3d as o3d
+import torch
+from attrdict import AttrDict
+from matplotlib.patches import Circle, Rectangle
+
+log = logging.getLogger(__name__)
+
+"""
+Open3d visualization functions
+"""
+
+
+class Visualizer3d:
+    def __init__(self, base_path="", view_params=None, tsleep=0.01):
+        self.vis = o3d.visualization.VisualizerWithKeyCallback()
+        self.vis.create_window(top=0, left=750, width=1080, height=1080)
+        self.tsleep = tsleep
+
+        self.base_path = base_path
+
+        self.view_params = AttrDict(
+            {
+                "fov": 0,
+                "front": [0.0, 0.0, 0.0],
+                "lookat": [0.0, 0.0, 0.0],
+                "up": [0.0, 0.0, 0.0],
+                "zoom": 0.5,
+            }
+        )
+        if view_params is not None:
+            self.view_params = view_params
+
+        self._init_options()
+
+    def _init_options(self):
+        opt = self.vis.get_render_option()
+        opt.show_coordinate_frame = True
+        opt.background_color = np.asarray([0.6, 0.6, 0.6])
+
+        # pause option
+        self.paused = types.SimpleNamespace()
+        self.paused.value = False
+        self.vis.register_key_action_callback(
+            ord("P"),
+            lambda a, b, c: b == 1
+            or setattr(self.paused, "value", not self.paused.value),
+        )
+
+    def _init_geom_cloud(self):
+        return o3d.geometry.PointCloud()
+
+    def _init_geom_frame(self, frame_size=0.01, frame_origin=[0, 0, 0]):
+        return o3d.geometry.TriangleMesh.create_coordinate_frame(
+            size=frame_size, origin=frame_origin
+        )
+
+    def _init_geom_mesh(self, mesh_name, color=None, wireframe=False):
+        mesh = o3d.io.read_triangle_mesh(
+            f"{self.base_path}/local/resources/meshes/{mesh_name}"
+        )
+        mesh.compute_vertex_normals()
+
+        if wireframe:
+            mesh = o3d.geometry.LineSet.create_from_triangle_mesh(mesh)
+        if color is not None:
+            mesh.paint_uniform_color(color)
+
+        return mesh
+
+    def init_geometry(
+        self,
+        geom_type,
+        num_items=1,
+        sizes=None,
+        file_names=None,
+        colors=None,
+        wireframes=None,
+    ):
+        geom_list = []
+        for i in range(0, num_items):
+
+            if geom_type == "cloud":
+                geom = self._init_geom_cloud()
+            elif geom_type == "frame":
+                frame_size = sizes[i] if sizes is not None else 0.001
+                geom = self._init_geom_frame(frame_size=frame_size)
+            elif geom_type == "mesh":
+                color = colors[i] if colors is not None else None
+                wireframe = wireframes[i] if wireframes is not None else False
+                geom = self._init_geom_mesh(file_names[i], color, wireframe)
+            else:
+                log.error(
+                    f"[Visualizer3d::init_geometry] geom_type {geom_type} not found."
+                )
+
+            geom_list.append(geom)
+
+        return geom_list
+
+    def add_geometry(self, geom_list):
+
+        if geom_list is None:
+            return
+
+        for geom in geom_list:
+            self.vis.add_geometry(geom)
+
+    def remove_geometry(self, geom_list, reset_bounding_box=False):
+
+        if geom_list is None:
+            return
+
+        for geom in geom_list:
+            self.vis.remove_geometry(geom, reset_bounding_box=reset_bounding_box)
+
+    def update_geometry(self, geom_list):
+        for geom in geom_list:
+            self.vis.update_geometry(geom)
+
+    def set_view(self):
+        ctr = self.vis.get_view_control()
+        ctr.change_field_of_view(self.view_params.fov)
+        ctr.set_front(self.view_params.front)
+        ctr.set_lookat(self.view_params.lookat)
+        ctr.set_up(self.view_params.up)
+        ctr.set_zoom(self.view_params.zoom)
+
+    def set_view_cam(self, T):
+        ctr = self.vis.get_view_control()
+        cam = ctr.convert_to_pinhole_camera_parameters()
+        cam.extrinsic = T
+        ctr.convert_from_pinhole_camera_parameters(cam)
+
+    def set_zoom(self):
+        ctr = self.vis.get_view_control()
+        ctr.set_zoom(1.5)
+
+    def rotate_view(self):
+        ctr = self.vis.get_view_control()
+        ctr.rotate(10.0, -0.0)
+
+    def pan_scene(self, max=300):
+        for i in range(0, max):
+            self.rotate_view()
+            self.render()
+
+    def render(self, T=None):
+
+        if T is not None:
+            self.set_view_cam(T)
+        else:
+            self.set_view()
+
+        self.vis.poll_events()
+        self.vis.update_renderer()
+        time.sleep(self.tsleep)
+
+    def transform_geometry_absolute(self, transform_list, geom_list):
+        for idx, geom in enumerate(geom_list):
+            T = transform_list[idx]
+            geom.transform(T)
+
+    def transform_geometry_relative(
+        self, transform_prev_list, transform_curr_list, geom_list
+    ):
+        for idx, geom in enumerate(geom_list):
+            T_prev = transform_prev_list[idx]
+            T_curr = transform_curr_list[idx]
+
+            # a. rotate R1^{-1}*R2 about center t1
+            geom.rotate(
+                torch.matmul(torch.inverse(T_prev[0:3, 0:3]), T_curr[0:3, 0:3]),
+                center=(T_prev[0, -1], T_prev[1, -1], T_prev[2, -1]),
+            )
+
+            # b. translate by t2 - t1
+            geom.translate(
+                (
+                    T_curr[0, -1] - T_prev[0, -1],
+                    T_curr[1, -1] - T_prev[1, -1],
+                    T_curr[2, -1] - T_prev[2, -1],
+                )
+            )
+
+    def clear_geometries(self):
+        self.vis.clear_geometries()
+
+    def destroy(self):
+        self.vis.destroy_window()
+
+
+def visualize_registration(source, target, transformation, vis3d=None, colors=None):
+    source_copy = copy.deepcopy(source)
+    target_copy = copy.deepcopy(target)
+
+    source_copy.transform(transformation)
+
+    clouds = [target_copy, source_copy]
+
+    if colors is not None:
+        clouds[0].paint_uniform_color(colors[1])
+        clouds[1].paint_uniform_color(colors[0])
+
+    vis3d.add_geometry(clouds)
+    vis3d.render()
+    vis3d.remove_geometry(clouds)
+
+
+def visualize_geometries_o3d(
+    vis3d, clouds=None, frames=None, meshes=None, transforms=None
+):
+    if meshes is not None:
+        meshes = [copy.deepcopy(mesh) for mesh in meshes]
+        if transforms is not None:
+            vis3d.transform_geometry_absolute(transforms, meshes)
+
+    if frames is not None:
+        frames = [copy.deepcopy(frame) for frame in frames]
+        if transforms is not None:
+            vis3d.transform_geometry_absolute(transforms, frames)
+
+    if clouds is not None:
+        vis3d.add_geometry(clouds)
+    if meshes is not None:
+        vis3d.add_geometry(meshes)
+    if frames is not None:
+        vis3d.add_geometry(frames)
+
+    vis3d.render()
+
+    if clouds is not None:
+        vis3d.remove_geometry(clouds)
+    if meshes is not None:
+        vis3d.remove_geometry(meshes)
+    if frames is not None:
+        vis3d.remove_geometry(frames)
+
+
+def visualize_inlier_outlier(cloud, ind):
+    inlier_cloud = cloud.select_by_index(ind)
+    outlier_cloud = cloud.select_by_index(ind, invert=True)
+
+    print("Showing outliers (red) and inliers (gray): ")
+    outlier_cloud.paint_uniform_color([1, 0, 0])
+    inlier_cloud.paint_uniform_color([0.8, 0.8, 0.8])
+    o3d.visualization.draw_geometries(
+        [inlier_cloud, outlier_cloud],
+        zoom=0.3412,
+        front=[0.4257, -0.2125, -0.8795],
+        lookat=[2.6172, 2.0475, 1.532],
+        up=[-0.0694, -0.9768, 0.2024],
+    )
+
+
+"""
+Optical flow visualization functions
+"""
+
+
+def flow_to_color(flow_uv, cvt=cv2.COLOR_HSV2BGR):
+    hsv = np.zeros((flow_uv.shape[0], flow_uv.shape[1], 3), dtype=np.uint8)
+    hsv[:, :, 0] = 255
+    hsv[:, :, 1] = 255
+    mag, ang = cv2.cartToPolar(flow_uv[..., 0], flow_uv[..., 1])
+    hsv[..., 0] = ang * 180 / np.pi / 2
+    hsv[..., 2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX)
+    flow_color = cv2.cvtColor(hsv, cvt)
+
+    return flow_color
+
+
+def flow_to_arrows(img, flow, step=8):
+    img = copy.deepcopy(img)
+    # img = (255 * img).astype(np.uint8)
+
+    h, w = img.shape[:2]
+    y, x = np.mgrid[step / 2 : h : step, step / 2 : w : step].reshape(2, -1).astype(int)
+    fx, fy = 5.0 * flow[y, x].T
+    lines = np.vstack([x, y, x + fx, y + fy]).T.reshape(-1, 2, 2)
+    lines = np.int32(lines + 0.5)
+    cv2.polylines(img, lines, 0, color=(0, 255, 0), thickness=1)
+    for (x1, y1), (x2, y2) in lines:
+        cv2.circle(img, (x1, y1), 1, (0, 255, 0), -1)
+
+    return img
+
+
+def depth_to_color(depth):
+    gray = (
+        np.clip((depth - depth.min()) / (depth.max() - depth.min()), 0, 1) * 255
+    ).astype(np.uint8)
+    return cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
+
+
+def visualize_flow_cv2(
+    img1, img2, flow_arrow=None, flow_color=None, win_size=(360, 360)
+):
+    img_disp1 = np.concatenate([img1, img2], axis=1)
+    cv2.namedWindow("img1, img2", cv2.WINDOW_NORMAL)
+    cv2.resizeWindow("img1, img2", 2 * win_size[0], win_size[1])
+    cv2.imshow("img1, img2", img_disp1)
+
+    if flow_arrow is not None:
+        cv2.namedWindow("flow_arrow", cv2.WINDOW_NORMAL)
+        cv2.resizeWindow("flow_arrow", win_size[0], win_size[1])
+        cv2.imshow("flow_arrow", flow_arrow)
+
+    if flow_color is not None:
+        cv2.namedWindow("flow_color", cv2.WINDOW_NORMAL)
+        cv2.resizeWindow("flow_color", win_size[0], win_size[1])
+        cv2.imshow("flow_color", flow_color)
+
+    cv2.waitKey(300)
+
+
+"""
+General visualization functions
+"""
+
+
+def draw_rectangle(
+    center_x,
+    center_y,
+    size_x,
+    size_y,
+    ang=0.0,
+    edgecolor="dimgray",
+    facecolor=None,
+    linewidth=2,
+):
+    R = np.array([[np.cos(ang), -np.sin(ang)], [np.sin(ang), np.cos(ang)]])
+    offset = np.matmul(R, np.array([[0.5 * size_x], [0.5 * size_y]]))
+    anchor_x = center_x - offset[0]
+    anchor_y = center_y - offset[1]
+    rect = Rectangle(
+        (anchor_x, anchor_y),
+        size_x,
+        size_y,
+        angle=(np.rad2deg(ang)),
+        facecolor=facecolor,
+        edgecolor=edgecolor,
+        linewidth=linewidth,
+    )
+    plt.gca().add_patch(rect)
+
+
+def draw_circle(center_x, center_y, radius):
+    circle = Circle((center_x, center_y), color="dimgray", radius=radius)
+    plt.gca().add_patch(circle)
+
+
+def visualize_imgs(fig, axs, img_list, titles=None, cmap=None):
+    for idx, img in enumerate(img_list):
+
+        if img is None:
+            continue
+
+        im = axs[idx].imshow(img, cmap=cmap)
+        if cmap is not None:
+            fig.colorbar(im, ax=axs[idx])
+        if titles is not None:
+            axs[idx].set_title(titles[idx])