wxchen/digit-depth
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

origin

Browse Source
This commit is contained in:
wxchen
2022-12-29 23:08:25 +08:00
commit 21ad625896
42 changed files with 2336 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

372
digit_depth/third_party/vis_utils.py vendored Normal file
View File

@@ -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])