Spaces:
Runtime error
Runtime error
File size: 6,762 Bytes
f670afc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 |
# 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
|