Spaces:
Build error
Build error
| """ | |
| Most of the code is taken from https://github.com/andyzeng/tsdf-fusion-python/blob/master/fusion.py | |
| @inproceedings{zeng20163dmatch, | |
| title={3DMatch: Learning Local Geometric Descriptors from RGB-D Reconstructions}, | |
| author={Zeng, Andy and Song, Shuran and Nie{\ss}ner, Matthias and Fisher, Matthew and Xiao, Jianxiong and Funkhouser, Thomas}, | |
| booktitle={CVPR}, | |
| year={2017} | |
| } | |
| """ | |
| import numpy as np | |
| from numba import njit, prange | |
| from skimage import measure | |
| FUSION_GPU_MODE = 0 | |
| class TSDFVolume: | |
| """Volumetric TSDF Fusion of RGB-D Images.""" | |
| def __init__(self, vol_bnds, voxel_size, use_gpu=True): | |
| """Constructor. | |
| Args: | |
| vol_bnds (ndarray): An ndarray of shape (3, 2). Specifies the | |
| xyz bounds (min/max) in meters. | |
| voxel_size (float): The volume discretization in meters. | |
| """ | |
| vol_bnds = np.asarray(vol_bnds) | |
| assert vol_bnds.shape == (3, 2), "[!] `vol_bnds` should be of shape (3, 2)." | |
| # Define voxel volume parameters | |
| self._vol_bnds = vol_bnds | |
| self._voxel_size = float(voxel_size) | |
| self._trunc_margin = 5 * self._voxel_size # truncation on SDF | |
| # self._trunc_margin = 10 # truncation on SDF | |
| self._color_const = 256 * 256 | |
| # Adjust volume bounds and ensure C-order contiguous | |
| self._vol_dim = ( | |
| np.ceil((self._vol_bnds[:, 1] - self._vol_bnds[:, 0]) / self._voxel_size) | |
| .copy(order="C") | |
| .astype(int) | |
| ) | |
| self._vol_bnds[:, 1] = self._vol_bnds[:, 0] + self._vol_dim * self._voxel_size | |
| self._vol_origin = self._vol_bnds[:, 0].copy(order="C").astype(np.float32) | |
| print( | |
| "Voxel volume size: {} x {} x {} - # points: {:,}".format( | |
| self._vol_dim[0], | |
| self._vol_dim[1], | |
| self._vol_dim[2], | |
| self._vol_dim[0] * self._vol_dim[1] * self._vol_dim[2], | |
| ) | |
| ) | |
| # Initialize pointers to voxel volume in CPU memory | |
| self._tsdf_vol_cpu = np.zeros(self._vol_dim).astype(np.float32) | |
| # for computing the cumulative moving average of observations per voxel | |
| self._weight_vol_cpu = np.zeros(self._vol_dim).astype(np.float32) | |
| self._color_vol_cpu = np.zeros(self._vol_dim).astype(np.float32) | |
| self.gpu_mode = use_gpu and FUSION_GPU_MODE | |
| # Copy voxel volumes to GPU | |
| if self.gpu_mode: | |
| self._tsdf_vol_gpu = cuda.mem_alloc(self._tsdf_vol_cpu.nbytes) | |
| cuda.memcpy_htod(self._tsdf_vol_gpu, self._tsdf_vol_cpu) | |
| self._weight_vol_gpu = cuda.mem_alloc(self._weight_vol_cpu.nbytes) | |
| cuda.memcpy_htod(self._weight_vol_gpu, self._weight_vol_cpu) | |
| self._color_vol_gpu = cuda.mem_alloc(self._color_vol_cpu.nbytes) | |
| cuda.memcpy_htod(self._color_vol_gpu, self._color_vol_cpu) | |
| # Cuda kernel function (C++) | |
| self._cuda_src_mod = SourceModule( | |
| """ | |
| __global__ void integrate(float * tsdf_vol, | |
| float * weight_vol, | |
| float * color_vol, | |
| float * vol_dim, | |
| float * vol_origin, | |
| float * cam_intr, | |
| float * cam_pose, | |
| float * other_params, | |
| float * color_im, | |
| float * depth_im) { | |
| // Get voxel index | |
| int gpu_loop_idx = (int) other_params[0]; | |
| int max_threads_per_block = blockDim.x; | |
| int block_idx = blockIdx.z*gridDim.y*gridDim.x+blockIdx.y*gridDim.x+blockIdx.x; | |
| int voxel_idx = gpu_loop_idx*gridDim.x*gridDim.y*gridDim.z*max_threads_per_block+block_idx*max_threads_per_block+threadIdx.x; | |
| int vol_dim_x = (int) vol_dim[0]; | |
| int vol_dim_y = (int) vol_dim[1]; | |
| int vol_dim_z = (int) vol_dim[2]; | |
| if (voxel_idx > vol_dim_x*vol_dim_y*vol_dim_z) | |
| return; | |
| // Get voxel grid coordinates (note: be careful when casting) | |
| float voxel_x = floorf(((float)voxel_idx)/((float)(vol_dim_y*vol_dim_z))); | |
| float voxel_y = floorf(((float)(voxel_idx-((int)voxel_x)*vol_dim_y*vol_dim_z))/((float)vol_dim_z)); | |
| float voxel_z = (float)(voxel_idx-((int)voxel_x)*vol_dim_y*vol_dim_z-((int)voxel_y)*vol_dim_z); | |
| // Voxel grid coordinates to world coordinates | |
| float voxel_size = other_params[1]; | |
| float pt_x = vol_origin[0]+voxel_x*voxel_size; | |
| float pt_y = vol_origin[1]+voxel_y*voxel_size; | |
| float pt_z = vol_origin[2]+voxel_z*voxel_size; | |
| // World coordinates to camera coordinates | |
| float tmp_pt_x = pt_x-cam_pose[0*4+3]; | |
| float tmp_pt_y = pt_y-cam_pose[1*4+3]; | |
| float tmp_pt_z = pt_z-cam_pose[2*4+3]; | |
| float cam_pt_x = cam_pose[0*4+0]*tmp_pt_x+cam_pose[1*4+0]*tmp_pt_y+cam_pose[2*4+0]*tmp_pt_z; | |
| float cam_pt_y = cam_pose[0*4+1]*tmp_pt_x+cam_pose[1*4+1]*tmp_pt_y+cam_pose[2*4+1]*tmp_pt_z; | |
| float cam_pt_z = cam_pose[0*4+2]*tmp_pt_x+cam_pose[1*4+2]*tmp_pt_y+cam_pose[2*4+2]*tmp_pt_z; | |
| // Camera coordinates to image pixels | |
| int pixel_x = (int) roundf(cam_intr[0*3+0]*(cam_pt_x/cam_pt_z)+cam_intr[0*3+2]); | |
| int pixel_y = (int) roundf(cam_intr[1*3+1]*(cam_pt_y/cam_pt_z)+cam_intr[1*3+2]); | |
| // Skip if outside view frustum | |
| int im_h = (int) other_params[2]; | |
| int im_w = (int) other_params[3]; | |
| if (pixel_x < 0 || pixel_x >= im_w || pixel_y < 0 || pixel_y >= im_h || cam_pt_z<0) | |
| return; | |
| // Skip invalid depth | |
| float depth_value = depth_im[pixel_y*im_w+pixel_x]; | |
| if (depth_value == 0) | |
| return; | |
| // Integrate TSDF | |
| float trunc_margin = other_params[4]; | |
| float depth_diff = depth_value-cam_pt_z; | |
| if (depth_diff < -trunc_margin) | |
| return; | |
| float dist = fmin(1.0f,depth_diff/trunc_margin); | |
| float w_old = weight_vol[voxel_idx]; | |
| float obs_weight = other_params[5]; | |
| float w_new = w_old + obs_weight; | |
| weight_vol[voxel_idx] = w_new; | |
| tsdf_vol[voxel_idx] = (tsdf_vol[voxel_idx]*w_old+obs_weight*dist)/w_new; | |
| // Integrate color | |
| float old_color = color_vol[voxel_idx]; | |
| float old_b = floorf(old_color/(256*256)); | |
| float old_g = floorf((old_color-old_b*256*256)/256); | |
| float old_r = old_color-old_b*256*256-old_g*256; | |
| float new_color = color_im[pixel_y*im_w+pixel_x]; | |
| float new_b = floorf(new_color/(256*256)); | |
| float new_g = floorf((new_color-new_b*256*256)/256); | |
| float new_r = new_color-new_b*256*256-new_g*256; | |
| new_b = fmin(roundf((old_b*w_old+obs_weight*new_b)/w_new),255.0f); | |
| new_g = fmin(roundf((old_g*w_old+obs_weight*new_g)/w_new),255.0f); | |
| new_r = fmin(roundf((old_r*w_old+obs_weight*new_r)/w_new),255.0f); | |
| color_vol[voxel_idx] = new_b*256*256+new_g*256+new_r; | |
| }""" | |
| ) | |
| self._cuda_integrate = self._cuda_src_mod.get_function("integrate") | |
| # Determine block/grid size on GPU | |
| gpu_dev = cuda.Device(0) | |
| self._max_gpu_threads_per_block = gpu_dev.MAX_THREADS_PER_BLOCK | |
| n_blocks = int( | |
| np.ceil( | |
| float(np.prod(self._vol_dim)) | |
| / float(self._max_gpu_threads_per_block) | |
| ) | |
| ) | |
| grid_dim_x = min(gpu_dev.MAX_GRID_DIM_X, int(np.floor(np.cbrt(n_blocks)))) | |
| grid_dim_y = min( | |
| gpu_dev.MAX_GRID_DIM_Y, int(np.floor(np.sqrt(n_blocks / grid_dim_x))) | |
| ) | |
| grid_dim_z = min( | |
| gpu_dev.MAX_GRID_DIM_Z, | |
| int(np.ceil(float(n_blocks) / float(grid_dim_x * grid_dim_y))), | |
| ) | |
| self._max_gpu_grid_dim = np.array( | |
| [grid_dim_x, grid_dim_y, grid_dim_z] | |
| ).astype(int) | |
| self._n_gpu_loops = int( | |
| np.ceil( | |
| float(np.prod(self._vol_dim)) | |
| / float( | |
| np.prod(self._max_gpu_grid_dim) | |
| * self._max_gpu_threads_per_block | |
| ) | |
| ) | |
| ) | |
| else: | |
| # Get voxel grid coordinates | |
| xv, yv, zv = np.meshgrid( | |
| range(self._vol_dim[0]), | |
| range(self._vol_dim[1]), | |
| range(self._vol_dim[2]), | |
| indexing="ij", | |
| ) | |
| self.vox_coords = ( | |
| np.concatenate( | |
| [xv.reshape(1, -1), yv.reshape(1, -1), zv.reshape(1, -1)], axis=0 | |
| ) | |
| .astype(int) | |
| .T | |
| ) | |
| def vox2world(vol_origin, vox_coords, vox_size, offsets=(0.5, 0.5, 0.5)): | |
| """Convert voxel grid coordinates to world coordinates.""" | |
| vol_origin = vol_origin.astype(np.float32) | |
| vox_coords = vox_coords.astype(np.float32) | |
| # print(np.min(vox_coords)) | |
| cam_pts = np.empty_like(vox_coords, dtype=np.float32) | |
| for i in prange(vox_coords.shape[0]): | |
| for j in range(3): | |
| cam_pts[i, j] = ( | |
| vol_origin[j] | |
| + (vox_size * vox_coords[i, j]) | |
| + vox_size * offsets[j] | |
| ) | |
| return cam_pts | |
| def cam2pix(cam_pts, intr): | |
| """Convert camera coordinates to pixel coordinates.""" | |
| intr = intr.astype(np.float32) | |
| fx, fy = intr[0, 0], intr[1, 1] | |
| cx, cy = intr[0, 2], intr[1, 2] | |
| pix = np.empty((cam_pts.shape[0], 2), dtype=np.int64) | |
| for i in prange(cam_pts.shape[0]): | |
| pix[i, 0] = int(np.round((cam_pts[i, 0] * fx / cam_pts[i, 2]) + cx)) | |
| pix[i, 1] = int(np.round((cam_pts[i, 1] * fy / cam_pts[i, 2]) + cy)) | |
| return pix | |
| def integrate_tsdf(tsdf_vol, dist, w_old, obs_weight): | |
| """Integrate the TSDF volume.""" | |
| tsdf_vol_int = np.empty_like(tsdf_vol, dtype=np.float32) | |
| # print(tsdf_vol.shape) | |
| w_new = np.empty_like(w_old, dtype=np.float32) | |
| for i in prange(len(tsdf_vol)): | |
| w_new[i] = w_old[i] + obs_weight | |
| tsdf_vol_int[i] = (w_old[i] * tsdf_vol[i] + obs_weight * dist[i]) / w_new[i] | |
| return tsdf_vol_int, w_new | |
| def integrate(self, color_im, depth_im, cam_intr, cam_pose, obs_weight=1.0): | |
| """Integrate an RGB-D frame into the TSDF volume. | |
| Args: | |
| color_im (ndarray): An RGB image of shape (H, W, 3). | |
| depth_im (ndarray): A depth image of shape (H, W). | |
| cam_intr (ndarray): The camera intrinsics matrix of shape (3, 3). | |
| cam_pose (ndarray): The camera pose (i.e. extrinsics) of shape (4, 4). | |
| obs_weight (float): The weight to assign for the current observation. A higher | |
| value | |
| """ | |
| im_h, im_w = depth_im.shape | |
| # Fold RGB color image into a single channel image | |
| color_im = color_im.astype(np.float32) | |
| color_im = np.floor( | |
| color_im[..., 2] * self._color_const | |
| + color_im[..., 1] * 256 | |
| + color_im[..., 0] | |
| ) | |
| if self.gpu_mode: # GPU mode: integrate voxel volume (calls CUDA kernel) | |
| for gpu_loop_idx in range(self._n_gpu_loops): | |
| self._cuda_integrate( | |
| self._tsdf_vol_gpu, | |
| self._weight_vol_gpu, | |
| self._color_vol_gpu, | |
| cuda.InOut(self._vol_dim.astype(np.float32)), | |
| cuda.InOut(self._vol_origin.astype(np.float32)), | |
| cuda.InOut(cam_intr.reshape(-1).astype(np.float32)), | |
| cuda.InOut(cam_pose.reshape(-1).astype(np.float32)), | |
| cuda.InOut( | |
| np.asarray( | |
| [ | |
| gpu_loop_idx, | |
| self._voxel_size, | |
| im_h, | |
| im_w, | |
| self._trunc_margin, | |
| obs_weight, | |
| ], | |
| np.float32, | |
| ) | |
| ), | |
| cuda.InOut(color_im.reshape(-1).astype(np.float32)), | |
| cuda.InOut(depth_im.reshape(-1).astype(np.float32)), | |
| block=(self._max_gpu_threads_per_block, 1, 1), | |
| grid=( | |
| int(self._max_gpu_grid_dim[0]), | |
| int(self._max_gpu_grid_dim[1]), | |
| int(self._max_gpu_grid_dim[2]), | |
| ), | |
| ) | |
| else: # CPU mode: integrate voxel volume (vectorized implementation) | |
| # Convert voxel grid coordinates to pixel coordinates | |
| cam_pts = self.vox2world( | |
| self._vol_origin, self.vox_coords, self._voxel_size | |
| ) | |
| cam_pts = rigid_transform(cam_pts, np.linalg.inv(cam_pose)) | |
| pix_z = cam_pts[:, 2] | |
| pix = self.cam2pix(cam_pts, cam_intr) | |
| pix_x, pix_y = pix[:, 0], pix[:, 1] | |
| # Eliminate pixels outside view frustum | |
| valid_pix = np.logical_and( | |
| pix_x >= 0, | |
| np.logical_and( | |
| pix_x < im_w, | |
| np.logical_and(pix_y >= 0, np.logical_and(pix_y < im_h, pix_z > 0)), | |
| ), | |
| ) | |
| depth_val = np.zeros(pix_x.shape) | |
| depth_val[valid_pix] = depth_im[pix_y[valid_pix], pix_x[valid_pix]] | |
| # Integrate TSDF | |
| depth_diff = depth_val - pix_z | |
| valid_pts = np.logical_and(depth_val > 0, depth_diff >= -10) | |
| dist = depth_diff | |
| valid_vox_x = self.vox_coords[valid_pts, 0] | |
| valid_vox_y = self.vox_coords[valid_pts, 1] | |
| valid_vox_z = self.vox_coords[valid_pts, 2] | |
| w_old = self._weight_vol_cpu[valid_vox_x, valid_vox_y, valid_vox_z] | |
| tsdf_vals = self._tsdf_vol_cpu[valid_vox_x, valid_vox_y, valid_vox_z] | |
| valid_dist = dist[valid_pts] | |
| tsdf_vol_new, w_new = self.integrate_tsdf( | |
| tsdf_vals, valid_dist, w_old, obs_weight | |
| ) | |
| self._weight_vol_cpu[valid_vox_x, valid_vox_y, valid_vox_z] = w_new | |
| self._tsdf_vol_cpu[valid_vox_x, valid_vox_y, valid_vox_z] = tsdf_vol_new | |
| # Integrate color | |
| old_color = self._color_vol_cpu[valid_vox_x, valid_vox_y, valid_vox_z] | |
| old_b = np.floor(old_color / self._color_const) | |
| old_g = np.floor((old_color - old_b * self._color_const) / 256) | |
| old_r = old_color - old_b * self._color_const - old_g * 256 | |
| new_color = color_im[pix_y[valid_pts], pix_x[valid_pts]] | |
| new_b = np.floor(new_color / self._color_const) | |
| new_g = np.floor((new_color - new_b * self._color_const) / 256) | |
| new_r = new_color - new_b * self._color_const - new_g * 256 | |
| new_b = np.minimum( | |
| 255.0, np.round((w_old * old_b + obs_weight * new_b) / w_new) | |
| ) | |
| new_g = np.minimum( | |
| 255.0, np.round((w_old * old_g + obs_weight * new_g) / w_new) | |
| ) | |
| new_r = np.minimum( | |
| 255.0, np.round((w_old * old_r + obs_weight * new_r) / w_new) | |
| ) | |
| self._color_vol_cpu[valid_vox_x, valid_vox_y, valid_vox_z] = ( | |
| new_b * self._color_const + new_g * 256 + new_r | |
| ) | |
| def get_volume(self): | |
| if self.gpu_mode: | |
| cuda.memcpy_dtoh(self._tsdf_vol_cpu, self._tsdf_vol_gpu) | |
| cuda.memcpy_dtoh(self._color_vol_cpu, self._color_vol_gpu) | |
| return self._tsdf_vol_cpu, self._color_vol_cpu | |
| def get_point_cloud(self): | |
| """Extract a point cloud from the voxel volume.""" | |
| tsdf_vol, color_vol = self.get_volume() | |
| # Marching cubes | |
| verts = measure.marching_cubes_lewiner(tsdf_vol, level=0)[0] | |
| verts_ind = np.round(verts).astype(int) | |
| verts = verts * self._voxel_size + self._vol_origin | |
| # Get vertex colors | |
| rgb_vals = color_vol[verts_ind[:, 0], verts_ind[:, 1], verts_ind[:, 2]] | |
| colors_b = np.floor(rgb_vals / self._color_const) | |
| colors_g = np.floor((rgb_vals - colors_b * self._color_const) / 256) | |
| colors_r = rgb_vals - colors_b * self._color_const - colors_g * 256 | |
| colors = np.floor(np.asarray([colors_r, colors_g, colors_b])).T | |
| colors = colors.astype(np.uint8) | |
| pc = np.hstack([verts, colors]) | |
| return pc | |
| def get_mesh(self): | |
| """Compute a mesh from the voxel volume using marching cubes.""" | |
| tsdf_vol, color_vol = self.get_volume() | |
| # Marching cubes | |
| verts, faces, norms, vals = measure.marching_cubes_lewiner(tsdf_vol, level=0) | |
| verts_ind = np.round(verts).astype(int) | |
| verts = ( | |
| verts * self._voxel_size + self._vol_origin | |
| ) # voxel grid coordinates to world coordinates | |
| # Get vertex colors | |
| rgb_vals = color_vol[verts_ind[:, 0], verts_ind[:, 1], verts_ind[:, 2]] | |
| colors_b = np.floor(rgb_vals / self._color_const) | |
| colors_g = np.floor((rgb_vals - colors_b * self._color_const) / 256) | |
| colors_r = rgb_vals - colors_b * self._color_const - colors_g * 256 | |
| colors = np.floor(np.asarray([colors_r, colors_g, colors_b])).T | |
| colors = colors.astype(np.uint8) | |
| return verts, faces, norms, colors | |
| def rigid_transform(xyz, transform): | |
| """Applies a rigid transform to an (N, 3) pointcloud.""" | |
| xyz_h = np.hstack([xyz, np.ones((len(xyz), 1), dtype=np.float32)]) | |
| xyz_t_h = np.dot(transform, xyz_h.T).T | |
| return xyz_t_h[:, :3] | |
| def get_view_frustum(depth_im, cam_intr, cam_pose): | |
| """Get corners of 3D camera view frustum of depth image""" | |
| im_h = depth_im.shape[0] | |
| im_w = depth_im.shape[1] | |
| max_depth = np.max(depth_im) | |
| view_frust_pts = np.array( | |
| [ | |
| (np.array([0, 0, 0, im_w, im_w]) - cam_intr[0, 2]) | |
| * np.array([0, max_depth, max_depth, max_depth, max_depth]) | |
| / cam_intr[0, 0], | |
| (np.array([0, 0, im_h, 0, im_h]) - cam_intr[1, 2]) | |
| * np.array([0, max_depth, max_depth, max_depth, max_depth]) | |
| / cam_intr[1, 1], | |
| np.array([0, max_depth, max_depth, max_depth, max_depth]), | |
| ] | |
| ) | |
| view_frust_pts = rigid_transform(view_frust_pts.T, cam_pose).T | |
| return view_frust_pts | |
| def meshwrite(filename, verts, faces, norms, colors): | |
| """Save a 3D mesh to a polygon .ply file.""" | |
| # Write header | |
| ply_file = open(filename, "w") | |
| ply_file.write("ply\n") | |
| ply_file.write("format ascii 1.0\n") | |
| ply_file.write("element vertex %d\n" % (verts.shape[0])) | |
| ply_file.write("property float x\n") | |
| ply_file.write("property float y\n") | |
| ply_file.write("property float z\n") | |
| ply_file.write("property float nx\n") | |
| ply_file.write("property float ny\n") | |
| ply_file.write("property float nz\n") | |
| ply_file.write("property uchar red\n") | |
| ply_file.write("property uchar green\n") | |
| ply_file.write("property uchar blue\n") | |
| ply_file.write("element face %d\n" % (faces.shape[0])) | |
| ply_file.write("property list uchar int vertex_index\n") | |
| ply_file.write("end_header\n") | |
| # Write vertex list | |
| for i in range(verts.shape[0]): | |
| ply_file.write( | |
| "%f %f %f %f %f %f %d %d %d\n" | |
| % ( | |
| verts[i, 0], | |
| verts[i, 1], | |
| verts[i, 2], | |
| norms[i, 0], | |
| norms[i, 1], | |
| norms[i, 2], | |
| colors[i, 0], | |
| colors[i, 1], | |
| colors[i, 2], | |
| ) | |
| ) | |
| # Write face list | |
| for i in range(faces.shape[0]): | |
| ply_file.write("3 %d %d %d\n" % (faces[i, 0], faces[i, 1], faces[i, 2])) | |
| ply_file.close() | |
| def pcwrite(filename, xyzrgb): | |
| """Save a point cloud to a polygon .ply file.""" | |
| xyz = xyzrgb[:, :3] | |
| rgb = xyzrgb[:, 3:].astype(np.uint8) | |
| # Write header | |
| ply_file = open(filename, "w") | |
| ply_file.write("ply\n") | |
| ply_file.write("format ascii 1.0\n") | |
| ply_file.write("element vertex %d\n" % (xyz.shape[0])) | |
| ply_file.write("property float x\n") | |
| ply_file.write("property float y\n") | |
| ply_file.write("property float z\n") | |
| ply_file.write("property uchar red\n") | |
| ply_file.write("property uchar green\n") | |
| ply_file.write("property uchar blue\n") | |
| ply_file.write("end_header\n") | |
| # Write vertex list | |
| for i in range(xyz.shape[0]): | |
| ply_file.write( | |
| "%f %f %f %d %d %d\n" | |
| % ( | |
| xyz[i, 0], | |
| xyz[i, 1], | |
| xyz[i, 2], | |
| rgb[i, 0], | |
| rgb[i, 1], | |
| rgb[i, 2], | |
| ) | |
| ) | |