# -------------------------------------------------------- # Visualization utilities. The code is adapted from Spann3r: # https://github.com/HengyiWang/spann3r/blob/main/spann3r/tools/vis.py # -------------------------------------------------------- import numpy as np from tqdm import tqdm import open3d as o3d import imageio import os import os.path as osp import matplotlib.pyplot as plt import matplotlib.colors as mcolors from skimage import exposure def render_scene(vis, geometry, camera_parameters, bg_color=[1,1,1], point_size=1., uint8=True): vis.clear_geometries() for g in geometry: vis.add_geometry(g) ctr = vis.get_view_control() ctr.convert_from_pinhole_camera_parameters(camera_parameters, allow_arbitrary=True) opt = vis.get_render_option() #调整点的大小 opt.point_size = point_size opt.background_color = np.array(bg_color) vis.poll_events() vis.update_renderer() image = vis.capture_screen_float_buffer(do_render=True) if not uint8: return image else: image_uint8 = (np.asarray(image) * 255).astype(np.uint8) return image_uint8 def render_frames(pts_all, image_all, camera_parameters, output_dir, mask=None, save_video=True, save_camera=True, init_ids=[], c2ws=None, vis_cam=False, save_stride=1, sample_ratio=1., incremental=True, save_name='render_frames', bg_color=[1, 1, 1], point_size=1., fps=10): t, h, w, _ = pts_all.shape vis = o3d.visualization.Visualizer() vis.create_window(width=960, height=544) render_frame_path = os.path.join(output_dir, save_name) os.makedirs(render_frame_path, exist_ok=True) if save_camera: o3d.io.write_pinhole_camera_parameters(os.path.join(render_frame_path, 'camera.json'), camera_parameters) video_path = os.path.join(output_dir, f'{save_name}.mp4') if save_video: writer = imageio.get_writer(video_path, fps=fps) # construct point cloud for initial window pcd = o3d.geometry.PointCloud() if init_ids is None: init_ids = [] if len(init_ids) > 0: init_ids = np.array(init_ids) init_masks = mask[init_ids] init_pts = pts_all[init_ids][init_masks] init_colors = image_all[init_ids][init_masks] if sample_ratio < 1.: sampled_idx = np.random.choice(len(init_pts), int(len(init_pts)*sample_ratio), replace=False) init_pts = init_pts[sampled_idx] init_colors = init_colors[sampled_idx] pcd.points = o3d.utility.Vector3dVector(init_pts) pcd.colors = o3d.utility.Vector3dVector(init_colors) vis.add_geometry(pcd) # visualize incremental reconstruction for i in tqdm(range(t), desc="Rendering incremental reconstruction"): if i not in init_ids: new_pts = pts_all[i].reshape(-1, 3) new_colors = image_all[i].reshape(-1, 3) if mask is not None: new_pts = new_pts[mask[i].reshape(-1)] new_colors = new_colors[mask[i].reshape(-1)] if sample_ratio < 1.: sampled_idx = np.random.choice(len(new_pts), int(len(new_pts)*sample_ratio), replace=False) new_pts = new_pts[sampled_idx] new_colors = new_colors[sampled_idx] if incremental: pcd.points.extend(o3d.utility.Vector3dVector(new_pts)) pcd.colors.extend(o3d.utility.Vector3dVector(new_colors)) else: pcd.points = o3d.utility.Vector3dVector(new_pts) pcd.colors = o3d.utility.Vector3dVector(new_colors) if (i+1) % save_stride != 0: continue geometry = [pcd] if vis_cam: geometry = geometry + draw_camera(c2ws[i], img=image_all[i]) image_uint8 = render_scene(vis, geometry, camera_parameters, bg_color=bg_color, point_size=point_size) frame_filename = f'frame_{i:03d}.png' imageio.imwrite(osp.join(render_frame_path, frame_filename), image_uint8) if save_video: writer.append_data(image_uint8) if save_video: writer.close() vis.destroy_window() def create_image_plane(img, c2w, scale=0.1): # simulate a image in 3D with point cloud H, W, _ = img.shape points = np.meshgrid(np.linspace(0, W-1, W), np.linspace(0, H-1, H)) points = np.stack(points, axis=-1).reshape(-1, 2) #translate the center of focal points -= np.array([W/2, H/2]) points *= 2*scale/W points = np.concatenate([points, 0.1*np.ones((len(points), 1))], axis=-1) colors = img.reshape(-1, 3) # no need for such resolution sample_stride = max(1, int(0.2/scale)) points = points[::sample_stride] colors = colors[::sample_stride] pcd = o3d.geometry.PointCloud() pcd.points = o3d.utility.Vector3dVector(points) pcd.colors = o3d.utility.Vector3dVector(colors) pcd.transform(c2w) return pcd def draw_camera(c2w, cam_width=0.2/2, cam_height=0.2/2, f=0.10, color=[0, 1, 0], show_axis=True, img=None): points = [[0, 0, 0], [-cam_width, -cam_height, f], [cam_width, -cam_height, f], [cam_width, cam_height, f], [-cam_width, cam_height, f]] lines = [[0, 1], [0, 2], [0, 3], [0, 4], [1, 2], [2, 3], [3, 4], [4, 1]] colors = [color for i in range(len(lines))] line_set = o3d.geometry.LineSet() line_set.points = o3d.utility.Vector3dVector(points) line_set.lines = o3d.utility.Vector2iVector(lines) line_set.colors = o3d.utility.Vector3dVector(colors) line_set.transform(c2w) res = [line_set] if show_axis: axis = o3d.geometry.TriangleMesh.create_coordinate_frame() axis.scale(min(cam_width, cam_height), np.array([0., 0., 0.])) axis.transform(c2w) res.append(axis) if img is not None: # draw image in the plane of the camera img_plane = create_image_plane(img, c2w) res.append(img_plane) return res def find_render_cam(pcd, poses_all=None, cam_width=0.016, cam_height=0.012, cam_f=0.02): last_camera_params = None def print_camera_pose(vis): nonlocal last_camera_params ctr = vis.get_view_control() camera_params = ctr.convert_to_pinhole_camera_parameters() last_camera_params = camera_params print("Intrinsic matrix:") print(camera_params.intrinsic.intrinsic_matrix) print("\nExtrinsic matrix:") print(camera_params.extrinsic) return False vis = o3d.visualization.VisualizerWithKeyCallback() vis.create_window(width=960, height=544) vis.add_geometry(pcd) if poses_all is not None: for pose in poses_all: for geometry in draw_camera(pose, cam_width, cam_height, cam_f): vis.add_geometry(geometry) opt = vis.get_render_option() opt.point_size = 1 opt.background_color = np.array([0, 0, 0]) print_camera_pose(vis) print("Press the space key to record the current rendering view.") vis.register_key_callback(32, print_camera_pose) while vis.poll_events(): vis.update_renderer() vis.destroy_window() return last_camera_params def vis_frame_preds(preds, type, save_path, norm_dims=(0, 1, 2), enhance_z=False, cmap=True, save_imgs=True, save_video=True, fps=10): if norm_dims is not None: min_val = preds.min(axis=norm_dims, keepdims=True) max_val = preds.max(axis=norm_dims, keepdims=True) preds = (preds - min_val) / (max_val - min_val) save_path = osp.join(save_path, type) if save_imgs: os.makedirs(save_path, exist_ok=True) if save_video: video_path = osp.join(osp.dirname(save_path), f'{type}.mp4') writer = imageio.get_writer(video_path, fps=fps) for frame_id in tqdm(range(preds.shape[0]), desc=f"Visualizing {type}"): pred_vis = preds[frame_id].astype(np.float32) if cmap: if preds.shape[-1] == 3: h = 1-pred_vis[...,0] s = 1-pred_vis[...,1] v = 1-pred_vis[...,2] if enhance_z: new_v = exposure.equalize_adapthist(v, clip_limit=0.01, nbins=256) v = new_v*0.2 + v*0.8 pred_vis = mcolors.hsv_to_rgb(np.stack([h, s, v], axis=-1)) elif len(pred_vis.shape)==2 or pred_vis.shape[-1] == 1: pred_vis = plt.cm.jet(pred_vis) pred_vis_rgb_uint8 = (pred_vis * 255).astype(np.uint8) if save_imgs: plt.imsave(osp.join(save_path, f'{type}_{frame_id:04d}.png'), pred_vis_rgb_uint8) if save_video: writer.append_data(pred_vis_rgb_uint8) if save_video: writer.close()