Spaces:
Runtime error
Runtime error
| # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # | |
| # This work is made available under the Nvidia Source Code License-NC. | |
| # To view a copy of this license, check out LICENSE.md | |
| import pickle | |
| import time | |
| import numpy as np | |
| class SplatRenderer(object): | |
| """Splatting 3D point cloud into image using precomputed mapping.""" | |
| def __init__(self): | |
| self.reset() | |
| def reset(self): | |
| """Reset the renderer.""" | |
| # 1 = point seen before, 0 = not seen. | |
| # This is numpy uint8 array of size (N, 1) | |
| self.seen_mask = None | |
| # Time of first colorization of 3D point. | |
| # This is numpy uint16 array of size (N, 1) | |
| self.seen_time = None | |
| # colors[kp_idx] is color of kp_idx'th keypoint. | |
| # This is a numpy uint8 array of size (N, 3) | |
| self.colors = None | |
| self.time_taken = 0 | |
| self.call_idx = 0 | |
| def num_points(self): | |
| r"""Number of points with assigned colors.""" | |
| return np.sum(self.seen_mask) | |
| def _resize_arrays(self, max_point_idx): | |
| r"""Makes arrays bigger, if needed. | |
| Args: | |
| max_point_idx (int): Highest 3D point index seen so far. | |
| """ | |
| if self.colors is None: | |
| old_max_point_idx = 0 | |
| else: | |
| old_max_point_idx = self.colors.shape[0] | |
| if max_point_idx > old_max_point_idx: | |
| # Init new bigger arrays. | |
| colors = np.zeros((max_point_idx, 3), dtype=np.uint8) | |
| seen_mask = np.zeros((max_point_idx, 1), dtype=np.uint8) | |
| seen_time = np.zeros((max_point_idx, 1), dtype=np.uint16) | |
| # Copy old colors, if exist. | |
| if old_max_point_idx > 0: | |
| colors[:old_max_point_idx] = self.colors | |
| seen_mask[:old_max_point_idx] = self.seen_mask | |
| seen_time[:old_max_point_idx] = self.seen_time | |
| # Reset pointers. | |
| self.colors = colors | |
| self.seen_mask = seen_mask | |
| self.seen_time = seen_time | |
| def update_point_cloud(self, image, point_info): | |
| r"""Updates point cloud with new points and colors. | |
| Args: | |
| image (H x W x 3, uint8): Select colors from this image to assign to | |
| 3D points which do not have previously assigned colors. | |
| point_info (N x 3): (i, j, 3D point idx) per row containing | |
| mapping of image pixel to 3D point in point cloud. | |
| """ | |
| if point_info is None or len(point_info) == 0: | |
| return | |
| start = time.time() | |
| self.call_idx += 1 | |
| i_idxs = point_info[:, 0] | |
| j_idxs = point_info[:, 1] | |
| point_idxs = point_info[:, 2] | |
| # Allocate memory for new colors. | |
| max_point_idx = np.max(np.array(point_idxs)) + 1 | |
| self._resize_arrays(max_point_idx) | |
| # print('max point idx:', max_point_idx) | |
| # Save only the new colors. | |
| self.colors[point_idxs] = \ | |
| self.seen_mask[point_idxs] * self.colors[point_idxs] + \ | |
| (1 - self.seen_mask[point_idxs]) * image[i_idxs, j_idxs] | |
| # Save point seen times. | |
| self.seen_time[point_idxs] = \ | |
| self.seen_mask[point_idxs] * self.seen_time[point_idxs] + \ | |
| (1 - self.seen_mask[point_idxs]) * self.call_idx | |
| # Update seen point mask. | |
| self.seen_mask[point_idxs] = 1 | |
| end = time.time() | |
| self.time_taken += (end - start) | |
| def render_image(self, point_info, w, h, return_mask=False): | |
| r"""Creates image of (h, w) and fills in colors. | |
| Args: | |
| point_info (N x 3): (i, j, 3D point idx) per row containing | |
| mapping of image pixel to 3D point in point cloud. | |
| w (int): Width of output image. | |
| h (int): Height of output image. | |
| return_mask (bool): Return binary mask of coloring. | |
| Returns: | |
| (tuple): | |
| - output (H x W x 3, uint8): Image formed with mapping and colors. | |
| - mask (H x W x 1, uint8): Binary (255 or 0) mask of colorization. | |
| """ | |
| output = np.zeros((h, w, 3), dtype=np.uint8) | |
| mask = np.zeros((h, w, 1), dtype=np.uint8) | |
| if point_info is None or len(point_info) == 0: | |
| if return_mask: | |
| return output, mask | |
| else: | |
| return output | |
| start = time.time() | |
| i_idxs = point_info[:, 0] | |
| j_idxs = point_info[:, 1] | |
| point_idxs = point_info[:, 2] | |
| # Allocate memory for new colors. | |
| max_point_idx = np.max(np.array(point_idxs)) + 1 | |
| self._resize_arrays(max_point_idx) | |
| # num_found = np.sum(self.seen_mask[point_idxs]) | |
| # print('Found %d points to color' % (num_found)) | |
| # Copy colors. | |
| output[i_idxs, j_idxs] = self.colors[point_idxs] | |
| end = time.time() | |
| self.time_taken += (end - start) | |
| if return_mask: | |
| mask[i_idxs, j_idxs] = 255 * self.seen_mask[point_idxs] | |
| return output, mask | |
| else: | |
| return output | |
| def decode_unprojections(data): | |
| r"""Unpickle unprojections and make array. | |
| Args: | |
| data (array of pickled info): Each pickled string has keypoint mapping | |
| info. | |
| Returns: | |
| output (dict): Keys are the different resolutions, and values are padded | |
| mapping information. | |
| """ | |
| # Unpickle unprojections and store them in a dict with resolutions as keys. | |
| all_unprojections = {} | |
| for item in data: | |
| info = pickle.loads(item) | |
| for resolution, value in info.items(): | |
| if resolution not in all_unprojections: | |
| all_unprojections[resolution] = [] | |
| if not value or value is None: | |
| point_info = [] | |
| else: | |
| point_info = value | |
| all_unprojections[resolution].append(point_info) | |
| outputs = {} | |
| for resolution, values in all_unprojections.items(): | |
| # Get max length of mapping. | |
| max_len = 0 | |
| for value in values: | |
| max_len = max(max_len, len(value)) | |
| # Entries are a 3-tuple of (i_idx, j_idx, point_idx). | |
| assert len(value) % 3 == 0 | |
| # Pad each mapping to max_len. | |
| values = [ | |
| value + # Original info. | |
| [-1] * (max_len - len(value)) + # Padding. | |
| [len(value) // 3] * 3 # End sentinel with length. | |
| for value in values | |
| ] | |
| # Convert each mapping to numpy and reshape. | |
| values = [np.array(value).reshape(-1, 3) for value in values] | |
| # Stack and put in output. | |
| # Shape is (T, N, 3). T is time steps, N is num mappings. | |
| outputs[resolution] = np.stack(values, axis=0) | |
| return outputs | |