373 lines
10 KiB
Python
373 lines
10 KiB
Python
# 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])
|