################################################################################################
### This implementation is taken and adapted from https://github.com/andyzeng/tsdf-fusion-python
### Copyright (c) 2018 Andy Zeng
################################################################################################

import time
from argparse import ArgumentParser

import cv2
import gc
import numpy as np
from numba import njit, prange
from path import Path
from skimage import measure

from dvmvs.dataset_loader import PreprocessImage, load_image

try:
    import pycuda.driver as cuda
    import pycuda.autoinit
    from pycuda.compiler import SourceModule

    FUSION_GPU_MODE = 1
except Exception as err:
    print('Warning: {}'.format(err))
    print('Failed to import PyCUDA. Running fusion in CPU mode.')
    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._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.ones(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

    @staticmethod
    @njit(parallel=True)
    def vox2world(vol_origin, vox_coords, vox_size):
        """Convert voxel grid coordinates to world coordinates.
        """
        vol_origin = vol_origin.astype(np.float32)
        vox_coords = vox_coords.astype(np.float32)
        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])
        return cam_pts

    @staticmethod
    @njit(parallel=True)
    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

    @staticmethod
    @njit(parallel=True)
    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)
        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.):
        """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 = TSDFFusion.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 >= -self._trunc_margin)
            dist = np.minimum(1, depth_diff / self._trunc_margin)
            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., np.round((w_old * old_b + obs_weight * new_b) / w_new))
            new_g = np.minimum(255., np.round((w_old * old_g + obs_weight * new_g) / w_new))
            new_r = np.minimum(255., 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


class TSDFFusion:
    @staticmethod
    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]

    @staticmethod
    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 = TSDFFusion.rigid_transform(view_frust_pts.T, cam_pose).T
        return view_frust_pts

    @staticmethod
    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()

    @staticmethod
    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],
            ))

    @staticmethod
    def integrate(tsdf_volume, images, depths, poses, K, mesh_name, save_progressive):
        n_imgs = len(images)
        # Loop through RGB-D images and fuse them together
        t0_elapse = time.time()
        for i in range(n_imgs):
            print("Fusing frame %d/%d" % (i + 1, n_imgs), "for", mesh_name)
            # Integrate observation into voxel volume (assume color aligned with depth)
            tsdf_volume.integrate(images[i], depths[i], K, poses[i], obs_weight=1.)

            if save_progressive:
                print("Saving for progressive visuals...")
                verts, faces, norms, colors = tsdf_volume.get_mesh()
                TSDFFusion.meshwrite(mesh_name + "_frame_{}.ply".format(str(i).zfill(5)), verts, faces, norms, colors)

        fps = n_imgs / (time.time() - t0_elapse)
        print("Average FPS: {:.2f}".format(fps))

        # Get mesh from voxel volume and save to disk (can be viewed with Meshlab)
        print("Saving mesh to", mesh_name)
        verts, faces, norms, colors = tsdf_volume.get_mesh()
        TSDFFusion.meshwrite(mesh_name + "_complete.ply", verts, faces, norms, colors)

    @staticmethod
    def calculate_volume_bounds(depth_maps, poses, K):
        assert len(depth_maps) == len(poses)

        volume_bounds = np.zeros((3, 2))
        for depth_map, pose in zip(depth_maps, poses):
            view_frust_points = TSDFFusion.get_view_frustum(depth_map, K, pose)
            volume_bounds[:, 0] = np.minimum(volume_bounds[:, 0], np.amin(view_frust_points, axis=1))
            volume_bounds[:, 1] = np.maximum(volume_bounds[:, 1], np.amax(view_frust_points, axis=1))

        return volume_bounds


def run(reconstruction_folder,
        prediction_folder,
        data_folder,
        dataset_name,
        scene_name,
        system_name,
        voxel_size,
        max_depth,
        use_groundtruth_to_anchor,
        save_progressive,
        save_groundtruth):
    #####
    ##### LIST RELEVANT FILENAMES #####
    #####
    dataset_folder = Path(data_folder) / dataset_name
    scene_folder = dataset_folder / scene_name

    original_K = np.loadtxt(scene_folder / 'K.txt').astype(np.float32)
    all_poses = np.fromfile(scene_folder / "poses.txt", dtype=float, sep="\n ").reshape((-1, 4, 4))
    all_image_filenames = sorted((scene_folder / 'images').files("*.png"))

    prediction_filename_regex = "keyframe_{}_{}_predictions_{}*".format(dataset_name, system_name, scene_name)
    prediction_filename = Path(prediction_folder).files(prediction_filename_regex)[0]
    predictions = np.load(prediction_filename)["arr_0"]
    prediction_height, prediction_width = np.shape(predictions[0])

    nmeas = system_name.split('_')[2]
    keyframe_filenames_regex = "keyframe+{}+{}+nmeas+{}".format(dataset_name, scene_name, nmeas)
    keyframe_filenames = np.loadtxt(Path(data_folder) / 'indices' / keyframe_filenames_regex, dtype=str, delimiter="\n")

    keyframe_poses = []
    keyframe_image_filenames = []
    for keyframe_filename in keyframe_filenames:
        if keyframe_filename == "TRACKING LOST":
            continue
        keyframe_filename = keyframe_filename.split(" ")[0]
        image_filename = scene_folder / 'images' / keyframe_filename
        pose_index = all_image_filenames.index(image_filename)
        keyframe_poses.append(all_poses[pose_index])
        keyframe_image_filenames.append(image_filename)

    #####
    ##### PREPARE IMAGES FOR TSDF RECONSTRUCTION #####
    #####
    temp_image = load_image(all_image_filenames[0])
    preprocessor = PreprocessImage(K=original_K,
                                   old_width=temp_image.shape[1],
                                   old_height=temp_image.shape[0],
                                   new_width=prediction_width,
                                   new_height=prediction_height,
                                   distortion_crop=0,
                                   perform_crop=False)
    scaled_K = preprocessor.get_updated_intrinsics()

    # ScanNet has black pixels around the edges of the image due to undistortion,
    # let's prepare a mask to ignore those corresponding predictions
    edge_pixel_amount = 10
    edge_mask = np.zeros((prediction_height, prediction_width), dtype=bool)
    edge_mask[0:edge_pixel_amount, :] = True
    edge_mask[prediction_height - edge_pixel_amount: prediction_height, :] = True
    edge_mask[:, 0:edge_pixel_amount] = True
    edge_mask[:, prediction_width - edge_pixel_amount: prediction_width] = True

    keyframe_images = []
    keyframe_predictions = []
    for index in range(0, len(keyframe_image_filenames)):
        keyframe_image = load_image(keyframe_image_filenames[index])
        keyframe_image = cv2.resize(keyframe_image, dsize=(prediction_width, prediction_height), interpolation=cv2.INTER_NEAREST)

        keyframe_prediction = predictions[index]

        if "scannet" in dataset_name:
            # masking the predictions corresponding to black pixels around the edges of the image
            black_mask = np.mean(keyframe_image.astype(float), axis=-1) < 10.0
            combined_mask = np.logical_and(black_mask, edge_mask)
            keyframe_prediction[combined_mask] = 0.0

        # masking the depth values larger than max_depth
        keyframe_prediction[keyframe_prediction > max_depth] = 0.0

        keyframe_predictions.append(keyframe_prediction)
        keyframe_images.append(keyframe_image.astype(np.uint8))

    #####
    ##### CALCULATE TSDF VOLUME BOUNDS WITH OR WITHOUT GROUNDTRUTH ANCHORING #####
    #####
    all_groundtruths = None
    if use_groundtruth_to_anchor or save_groundtruth:
        all_groundtruth_filenames = sorted((scene_folder / 'depth').files("*.png"))
        all_groundtruths = []
        for groundtruth_filename in all_groundtruth_filenames:
            groundtruth = cv2.imread(groundtruth_filename, -1).astype(float) / 1000.0
            groundtruth[groundtruth > max_depth] = 0.0
            all_groundtruths.append(groundtruth)

    if use_groundtruth_to_anchor:
        volume_bounds = TSDFFusion.calculate_volume_bounds(all_groundtruths, all_poses, original_K)
    else:
        volume_bounds = TSDFFusion.calculate_volume_bounds(keyframe_predictions, keyframe_poses, scaled_K)
    volume_bounds *= 1.05  # give some margin for the volume bounds accounting for mistakes

    #####
    ##### RUN THE RECONSTRUCTION WITH GROUNDTRUTH DEPTH MAPS #####
    #####
    if save_groundtruth:
        tsdf_volume = TSDFVolume(volume_bounds, voxel_size=voxel_size)

        mesh_name = "{}/reconstruction_voxelsize-{}_maxdepth-{}_anchor-{}_{}_{}_{}".format(reconstruction_folder,
                                                                                          voxel_size,
                                                                                          max_depth,
                                                                                          use_groundtruth_to_anchor,
                                                                                          "GROUNDTRUTH",
                                                                                          dataset_name,
                                                                                          scene_name)
        all_images = []
        for image_filename in all_image_filenames:
            image = load_image(image_filename)
            all_images.append(image)

        TSDFFusion.integrate(tsdf_volume=tsdf_volume,
                             images=all_images,
                             depths=all_groundtruths,
                             poses=all_poses,
                             K=original_K,
                             mesh_name=mesh_name,
                             save_progressive=False)

        tsdf_volume = None
        gc.collect()

    #####
    ##### RUN THE RECONSTRUCTION WITH PREDICTED DEPTH MAPS #####
    #####
    tsdf_volume = TSDFVolume(volume_bounds, voxel_size=voxel_size)

    mesh_name = "{}/reconstruction_voxelsize-{}_maxdepth-{}_anchor-{}_{}_{}_{}".format(reconstruction_folder,
                                                                                       voxel_size,
                                                                                       max_depth,
                                                                                       use_groundtruth_to_anchor,
                                                                                       system_name,
                                                                                       dataset_name,
                                                                                       scene_name)

    TSDFFusion.integrate(tsdf_volume=tsdf_volume,
                         images=keyframe_images,
                         depths=keyframe_predictions,
                         poses=keyframe_poses,
                         K=scaled_K,
                         mesh_name=mesh_name,
                         save_progressive=save_progressive)


if __name__ == "__main__":
    parser = ArgumentParser()

    parser.add_argument('--reconstruction_folder', default="./reconstructions", type=str)
    parser.add_argument('--prediction_folder', default="./predictions", type=str)
    parser.add_argument('--data_folder', default=".", type=str)
    parser.add_argument('--dataset_name', default="hololens-dataset", type=str)
    parser.add_argument('--scene_name', default="000", type=str)
    parser.add_argument('--system_name', default='320_256_3_dvmvs_fusionnet_online', type=str)
    parser.add_argument('--voxel_size', default=0.025, type=float)
    parser.add_argument('--max_depth', default=5.0, type=float)

    parser.add_argument('--use_groundtruth_to_anchor', action='store_true',
                        help="flag for using groundtruth depth maps to calculate the volume bounds and the origin of the volume,"
                             "recommended if groundtruth depth maps are available")
    parser.add_argument('--save_progressive', action='store_true',
                        help="flag for saving the mesh after each fused keyframe to the disk to get the progressive visuals "
                             "similar to the demo video shown in README")
    parser.add_argument('--save_groundtruth', action='store_true',
                        help="flag for saving the complete reconstruction resulting from groundtruth depth maps")

    args = parser.parse_args()

    run(reconstruction_folder=args.reconstruction_folder,
        prediction_folder=args.prediction_folder,
        data_folder=args.data_folder,
        dataset_name=args.dataset_name,
        scene_name=args.scene_name,
        system_name=args.system_name,
        voxel_size=args.voxel_size,
        max_depth=args.max_depth,
        use_groundtruth_to_anchor=args.use_groundtruth_to_anchor,
        save_progressive=args.save_progressive,
        save_groundtruth=args.save_groundtruth)

390e8b60a292d88f2bbae1f1eec11bb0636235db
<mask token>


class TSDFVolume:
    <mask token>

    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).'
        self._vol_bnds = vol_bnds
        self._voxel_size = float(voxel_size)
        self._trunc_margin = 5 * self._voxel_size
        self._color_const = 256 * 256
        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]))
        self._tsdf_vol_cpu = np.ones(self._vol_dim).astype(np.float32)
        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
        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)
            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')
            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:
            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
    <mask token>

    @staticmethod
    @njit(parallel=True)
    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

    @staticmethod
    @njit(parallel=True)
    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)
        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
        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:
            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:
            cam_pts = self.vox2world(self._vol_origin, self.vox_coords,
                self._voxel_size)
            cam_pts = TSDFFusion.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]
            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]]
            depth_diff = depth_val - pix_z
            valid_pts = np.logical_and(depth_val > 0, depth_diff >= -self.
                _trunc_margin)
            dist = np.minimum(1, depth_diff / self._trunc_margin)
            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
            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
    <mask token>

    def get_point_cloud(self):
        """Extract a point cloud from the voxel volume.
        """
        tsdf_vol, color_vol = self.get_volume()
        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
        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()
        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
        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


class TSDFFusion:

    @staticmethod
    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]

    @staticmethod
    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 = TSDFFusion.rigid_transform(view_frust_pts.T, cam_pose
            ).T
        return view_frust_pts

    @staticmethod
    def meshwrite(filename, verts, faces, norms, colors):
        """Save a 3D mesh to a polygon .ply file.
        """
        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')
        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]))
        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()

    @staticmethod
    def pcwrite(filename, xyzrgb):
        """Save a point cloud to a polygon .ply file.
        """
        xyz = xyzrgb[:, :3]
        rgb = xyzrgb[:, 3:].astype(np.uint8)
        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')
        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]))

    @staticmethod
    def integrate(tsdf_volume, images, depths, poses, K, mesh_name,
        save_progressive):
        n_imgs = len(images)
        t0_elapse = time.time()
        for i in range(n_imgs):
            print('Fusing frame %d/%d' % (i + 1, n_imgs), 'for', mesh_name)
            tsdf_volume.integrate(images[i], depths[i], K, poses[i],
                obs_weight=1.0)
            if save_progressive:
                print('Saving for progressive visuals...')
                verts, faces, norms, colors = tsdf_volume.get_mesh()
                TSDFFusion.meshwrite(mesh_name + '_frame_{}.ply'.format(str
                    (i).zfill(5)), verts, faces, norms, colors)
        fps = n_imgs / (time.time() - t0_elapse)
        print('Average FPS: {:.2f}'.format(fps))
        print('Saving mesh to', mesh_name)
        verts, faces, norms, colors = tsdf_volume.get_mesh()
        TSDFFusion.meshwrite(mesh_name + '_complete.ply', verts, faces,
            norms, colors)

    @staticmethod
    def calculate_volume_bounds(depth_maps, poses, K):
        assert len(depth_maps) == len(poses)
        volume_bounds = np.zeros((3, 2))
        for depth_map, pose in zip(depth_maps, poses):
            view_frust_points = TSDFFusion.get_view_frustum(depth_map, K, pose)
            volume_bounds[:, 0] = np.minimum(volume_bounds[:, 0], np.amin(
                view_frust_points, axis=1))
            volume_bounds[:, 1] = np.maximum(volume_bounds[:, 1], np.amax(
                view_frust_points, axis=1))
        return volume_bounds


<mask token>

<mask token>


class TSDFVolume:
    <mask token>

    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).'
        self._vol_bnds = vol_bnds
        self._voxel_size = float(voxel_size)
        self._trunc_margin = 5 * self._voxel_size
        self._color_const = 256 * 256
        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]))
        self._tsdf_vol_cpu = np.ones(self._vol_dim).astype(np.float32)
        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
        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)
            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')
            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:
            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

    @staticmethod
    @njit(parallel=True)
    def vox2world(vol_origin, vox_coords, vox_size):
        """Convert voxel grid coordinates to world coordinates.
        """
        vol_origin = vol_origin.astype(np.float32)
        vox_coords = vox_coords.astype(np.float32)
        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]
        return cam_pts

    @staticmethod
    @njit(parallel=True)
    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

    @staticmethod
    @njit(parallel=True)
    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)
        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
        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:
            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:
            cam_pts = self.vox2world(self._vol_origin, self.vox_coords,
                self._voxel_size)
            cam_pts = TSDFFusion.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]
            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]]
            depth_diff = depth_val - pix_z
            valid_pts = np.logical_and(depth_val > 0, depth_diff >= -self.
                _trunc_margin)
            dist = np.minimum(1, depth_diff / self._trunc_margin)
            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
            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
    <mask token>

    def get_point_cloud(self):
        """Extract a point cloud from the voxel volume.
        """
        tsdf_vol, color_vol = self.get_volume()
        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
        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()
        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
        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


class TSDFFusion:

    @staticmethod
    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]

    @staticmethod
    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 = TSDFFusion.rigid_transform(view_frust_pts.T, cam_pose
            ).T
        return view_frust_pts

    @staticmethod
    def meshwrite(filename, verts, faces, norms, colors):
        """Save a 3D mesh to a polygon .ply file.
        """
        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')
        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]))
        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()

    @staticmethod
    def pcwrite(filename, xyzrgb):
        """Save a point cloud to a polygon .ply file.
        """
        xyz = xyzrgb[:, :3]
        rgb = xyzrgb[:, 3:].astype(np.uint8)
        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')
        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]))

    @staticmethod
    def integrate(tsdf_volume, images, depths, poses, K, mesh_name,
        save_progressive):
        n_imgs = len(images)
        t0_elapse = time.time()
        for i in range(n_imgs):
            print('Fusing frame %d/%d' % (i + 1, n_imgs), 'for', mesh_name)
            tsdf_volume.integrate(images[i], depths[i], K, poses[i],
                obs_weight=1.0)
            if save_progressive:
                print('Saving for progressive visuals...')
                verts, faces, norms, colors = tsdf_volume.get_mesh()
                TSDFFusion.meshwrite(mesh_name + '_frame_{}.ply'.format(str
                    (i).zfill(5)), verts, faces, norms, colors)
        fps = n_imgs / (time.time() - t0_elapse)
        print('Average FPS: {:.2f}'.format(fps))
        print('Saving mesh to', mesh_name)
        verts, faces, norms, colors = tsdf_volume.get_mesh()
        TSDFFusion.meshwrite(mesh_name + '_complete.ply', verts, faces,
            norms, colors)

    @staticmethod
    def calculate_volume_bounds(depth_maps, poses, K):
        assert len(depth_maps) == len(poses)
        volume_bounds = np.zeros((3, 2))
        for depth_map, pose in zip(depth_maps, poses):
            view_frust_points = TSDFFusion.get_view_frustum(depth_map, K, pose)
            volume_bounds[:, 0] = np.minimum(volume_bounds[:, 0], np.amin(
                view_frust_points, axis=1))
            volume_bounds[:, 1] = np.maximum(volume_bounds[:, 1], np.amax(
                view_frust_points, axis=1))
        return volume_bounds


<mask token>

<mask token>


class TSDFVolume:
    <mask token>

    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).'
        self._vol_bnds = vol_bnds
        self._voxel_size = float(voxel_size)
        self._trunc_margin = 5 * self._voxel_size
        self._color_const = 256 * 256
        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]))
        self._tsdf_vol_cpu = np.ones(self._vol_dim).astype(np.float32)
        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
        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)
            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')
            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:
            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

    @staticmethod
    @njit(parallel=True)
    def vox2world(vol_origin, vox_coords, vox_size):
        """Convert voxel grid coordinates to world coordinates.
        """
        vol_origin = vol_origin.astype(np.float32)
        vox_coords = vox_coords.astype(np.float32)
        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]
        return cam_pts

    @staticmethod
    @njit(parallel=True)
    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

    @staticmethod
    @njit(parallel=True)
    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)
        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
        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:
            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:
            cam_pts = self.vox2world(self._vol_origin, self.vox_coords,
                self._voxel_size)
            cam_pts = TSDFFusion.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]
            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]]
            depth_diff = depth_val - pix_z
            valid_pts = np.logical_and(depth_val > 0, depth_diff >= -self.
                _trunc_margin)
            dist = np.minimum(1, depth_diff / self._trunc_margin)
            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
            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()
        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
        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()
        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
        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


class TSDFFusion:

    @staticmethod
    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]

    @staticmethod
    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 = TSDFFusion.rigid_transform(view_frust_pts.T, cam_pose
            ).T
        return view_frust_pts

    @staticmethod
    def meshwrite(filename, verts, faces, norms, colors):
        """Save a 3D mesh to a polygon .ply file.
        """
        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')
        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]))
        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()

    @staticmethod
    def pcwrite(filename, xyzrgb):
        """Save a point cloud to a polygon .ply file.
        """
        xyz = xyzrgb[:, :3]
        rgb = xyzrgb[:, 3:].astype(np.uint8)
        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')
        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]))

    @staticmethod
    def integrate(tsdf_volume, images, depths, poses, K, mesh_name,
        save_progressive):
        n_imgs = len(images)
        t0_elapse = time.time()
        for i in range(n_imgs):
            print('Fusing frame %d/%d' % (i + 1, n_imgs), 'for', mesh_name)
            tsdf_volume.integrate(images[i], depths[i], K, poses[i],
                obs_weight=1.0)
            if save_progressive:
                print('Saving for progressive visuals...')
                verts, faces, norms, colors = tsdf_volume.get_mesh()
                TSDFFusion.meshwrite(mesh_name + '_frame_{}.ply'.format(str
                    (i).zfill(5)), verts, faces, norms, colors)
        fps = n_imgs / (time.time() - t0_elapse)
        print('Average FPS: {:.2f}'.format(fps))
        print('Saving mesh to', mesh_name)
        verts, faces, norms, colors = tsdf_volume.get_mesh()
        TSDFFusion.meshwrite(mesh_name + '_complete.ply', verts, faces,
            norms, colors)

    @staticmethod
    def calculate_volume_bounds(depth_maps, poses, K):
        assert len(depth_maps) == len(poses)
        volume_bounds = np.zeros((3, 2))
        for depth_map, pose in zip(depth_maps, poses):
            view_frust_points = TSDFFusion.get_view_frustum(depth_map, K, pose)
            volume_bounds[:, 0] = np.minimum(volume_bounds[:, 0], np.amin(
                view_frust_points, axis=1))
            volume_bounds[:, 1] = np.maximum(volume_bounds[:, 1], np.amax(
                view_frust_points, axis=1))
        return volume_bounds


<mask token>

<mask token>


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).'
        self._vol_bnds = vol_bnds
        self._voxel_size = float(voxel_size)
        self._trunc_margin = 5 * self._voxel_size
        self._color_const = 256 * 256
        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]))
        self._tsdf_vol_cpu = np.ones(self._vol_dim).astype(np.float32)
        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
        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)
            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')
            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:
            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

    @staticmethod
    @njit(parallel=True)
    def vox2world(vol_origin, vox_coords, vox_size):
        """Convert voxel grid coordinates to world coordinates.
        """
        vol_origin = vol_origin.astype(np.float32)
        vox_coords = vox_coords.astype(np.float32)
        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]
        return cam_pts

    @staticmethod
    @njit(parallel=True)
    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

    @staticmethod
    @njit(parallel=True)
    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)
        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
        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:
            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:
            cam_pts = self.vox2world(self._vol_origin, self.vox_coords,
                self._voxel_size)
            cam_pts = TSDFFusion.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]
            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]]
            depth_diff = depth_val - pix_z
            valid_pts = np.logical_and(depth_val > 0, depth_diff >= -self.
                _trunc_margin)
            dist = np.minimum(1, depth_diff / self._trunc_margin)
            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
            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()
        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
        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()
        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
        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


class TSDFFusion:

    @staticmethod
    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]

    @staticmethod
    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 = TSDFFusion.rigid_transform(view_frust_pts.T, cam_pose
            ).T
        return view_frust_pts

    @staticmethod
    def meshwrite(filename, verts, faces, norms, colors):
        """Save a 3D mesh to a polygon .ply file.
        """
        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')
        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]))
        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()

    @staticmethod
    def pcwrite(filename, xyzrgb):
        """Save a point cloud to a polygon .ply file.
        """
        xyz = xyzrgb[:, :3]
        rgb = xyzrgb[:, 3:].astype(np.uint8)
        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')
        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]))

    @staticmethod
    def integrate(tsdf_volume, images, depths, poses, K, mesh_name,
        save_progressive):
        n_imgs = len(images)
        t0_elapse = time.time()
        for i in range(n_imgs):
            print('Fusing frame %d/%d' % (i + 1, n_imgs), 'for', mesh_name)
            tsdf_volume.integrate(images[i], depths[i], K, poses[i],
                obs_weight=1.0)
            if save_progressive:
                print('Saving for progressive visuals...')
                verts, faces, norms, colors = tsdf_volume.get_mesh()
                TSDFFusion.meshwrite(mesh_name + '_frame_{}.ply'.format(str
                    (i).zfill(5)), verts, faces, norms, colors)
        fps = n_imgs / (time.time() - t0_elapse)
        print('Average FPS: {:.2f}'.format(fps))
        print('Saving mesh to', mesh_name)
        verts, faces, norms, colors = tsdf_volume.get_mesh()
        TSDFFusion.meshwrite(mesh_name + '_complete.ply', verts, faces,
            norms, colors)

    @staticmethod
    def calculate_volume_bounds(depth_maps, poses, K):
        assert len(depth_maps) == len(poses)
        volume_bounds = np.zeros((3, 2))
        for depth_map, pose in zip(depth_maps, poses):
            view_frust_points = TSDFFusion.get_view_frustum(depth_map, K, pose)
            volume_bounds[:, 0] = np.minimum(volume_bounds[:, 0], np.amin(
                view_frust_points, axis=1))
            volume_bounds[:, 1] = np.maximum(volume_bounds[:, 1], np.amax(
                view_frust_points, axis=1))
        return volume_bounds


<mask token>

################################################################################################
### This implementation is taken and adapted from https://github.com/andyzeng/tsdf-fusion-python
### Copyright (c) 2018 Andy Zeng
################################################################################################

import time
from argparse import ArgumentParser

import cv2
import gc
import numpy as np
from numba import njit, prange
from path import Path
from skimage import measure

from dvmvs.dataset_loader import PreprocessImage, load_image

try:
    import pycuda.driver as cuda
    import pycuda.autoinit
    from pycuda.compiler import SourceModule

    FUSION_GPU_MODE = 1
except Exception as err:
    print('Warning: {}'.format(err))
    print('Failed to import PyCUDA. Running fusion in CPU mode.')
    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._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.ones(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

    @staticmethod
    @njit(parallel=True)
    def vox2world(vol_origin, vox_coords, vox_size):
        """Convert voxel grid coordinates to world coordinates.
        """
        vol_origin = vol_origin.astype(np.float32)
        vox_coords = vox_coords.astype(np.float32)
        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])
        return cam_pts

    @staticmethod
    @njit(parallel=True)
    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

    @staticmethod
    @njit(parallel=True)
    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)
        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.):
        """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 = TSDFFusion.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 >= -self._trunc_margin)
            dist = np.minimum(1, depth_diff / self._trunc_margin)
            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., np.round((w_old * old_b + obs_weight * new_b) / w_new))
            new_g = np.minimum(255., np.round((w_old * old_g + obs_weight * new_g) / w_new))
            new_r = np.minimum(255., 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


class TSDFFusion:
    @staticmethod
    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]

    @staticmethod
    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 = TSDFFusion.rigid_transform(view_frust_pts.T, cam_pose).T
        return view_frust_pts

    @staticmethod
    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()

    @staticmethod
    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],
            ))

    @staticmethod
    def integrate(tsdf_volume, images, depths, poses, K, mesh_name, save_progressive):
        n_imgs = len(images)
        # Loop through RGB-D images and fuse them together
        t0_elapse = time.time()
        for i in range(n_imgs):
            print("Fusing frame %d/%d" % (i + 1, n_imgs), "for", mesh_name)
            # Integrate observation into voxel volume (assume color aligned with depth)
            tsdf_volume.integrate(images[i], depths[i], K, poses[i], obs_weight=1.)

            if save_progressive:
                print("Saving for progressive visuals...")
                verts, faces, norms, colors = tsdf_volume.get_mesh()
                TSDFFusion.meshwrite(mesh_name + "_frame_{}.ply".format(str(i).zfill(5)), verts, faces, norms, colors)

        fps = n_imgs / (time.time() - t0_elapse)
        print("Average FPS: {:.2f}".format(fps))

        # Get mesh from voxel volume and save to disk (can be viewed with Meshlab)
        print("Saving mesh to", mesh_name)
        verts, faces, norms, colors = tsdf_volume.get_mesh()
        TSDFFusion.meshwrite(mesh_name + "_complete.ply", verts, faces, norms, colors)

    @staticmethod
    def calculate_volume_bounds(depth_maps, poses, K):
        assert len(depth_maps) == len(poses)

        volume_bounds = np.zeros((3, 2))
        for depth_map, pose in zip(depth_maps, poses):
            view_frust_points = TSDFFusion.get_view_frustum(depth_map, K, pose)
            volume_bounds[:, 0] = np.minimum(volume_bounds[:, 0], np.amin(view_frust_points, axis=1))
            volume_bounds[:, 1] = np.maximum(volume_bounds[:, 1], np.amax(view_frust_points, axis=1))

        return volume_bounds


def run(reconstruction_folder,
        prediction_folder,
        data_folder,
        dataset_name,
        scene_name,
        system_name,
        voxel_size,
        max_depth,
        use_groundtruth_to_anchor,
        save_progressive,
        save_groundtruth):
    #####
    ##### LIST RELEVANT FILENAMES #####
    #####
    dataset_folder = Path(data_folder) / dataset_name
    scene_folder = dataset_folder / scene_name

    original_K = np.loadtxt(scene_folder / 'K.txt').astype(np.float32)
    all_poses = np.fromfile(scene_folder / "poses.txt", dtype=float, sep="\n ").reshape((-1, 4, 4))
    all_image_filenames = sorted((scene_folder / 'images').files("*.png"))

    prediction_filename_regex = "keyframe_{}_{}_predictions_{}*".format(dataset_name, system_name, scene_name)
    prediction_filename = Path(prediction_folder).files(prediction_filename_regex)[0]
    predictions = np.load(prediction_filename)["arr_0"]
    prediction_height, prediction_width = np.shape(predictions[0])

    nmeas = system_name.split('_')[2]
    keyframe_filenames_regex = "keyframe+{}+{}+nmeas+{}".format(dataset_name, scene_name, nmeas)
    keyframe_filenames = np.loadtxt(Path(data_folder) / 'indices' / keyframe_filenames_regex, dtype=str, delimiter="\n")

    keyframe_poses = []
    keyframe_image_filenames = []
    for keyframe_filename in keyframe_filenames:
        if keyframe_filename == "TRACKING LOST":
            continue
        keyframe_filename = keyframe_filename.split(" ")[0]
        image_filename = scene_folder / 'images' / keyframe_filename
        pose_index = all_image_filenames.index(image_filename)
        keyframe_poses.append(all_poses[pose_index])
        keyframe_image_filenames.append(image_filename)

    #####
    ##### PREPARE IMAGES FOR TSDF RECONSTRUCTION #####
    #####
    temp_image = load_image(all_image_filenames[0])
    preprocessor = PreprocessImage(K=original_K,
                                   old_width=temp_image.shape[1],
                                   old_height=temp_image.shape[0],
                                   new_width=prediction_width,
                                   new_height=prediction_height,
                                   distortion_crop=0,
                                   perform_crop=False)
    scaled_K = preprocessor.get_updated_intrinsics()

    # ScanNet has black pixels around the edges of the image due to undistortion,
    # let's prepare a mask to ignore those corresponding predictions
    edge_pixel_amount = 10
    edge_mask = np.zeros((prediction_height, prediction_width), dtype=bool)
    edge_mask[0:edge_pixel_amount, :] = True
    edge_mask[prediction_height - edge_pixel_amount: prediction_height, :] = True
    edge_mask[:, 0:edge_pixel_amount] = True
    edge_mask[:, prediction_width - edge_pixel_amount: prediction_width] = True

    keyframe_images = []
    keyframe_predictions = []
    for index in range(0, len(keyframe_image_filenames)):
        keyframe_image = load_image(keyframe_image_filenames[index])
        keyframe_image = cv2.resize(keyframe_image, dsize=(prediction_width, prediction_height), interpolation=cv2.INTER_NEAREST)

        keyframe_prediction = predictions[index]

        if "scannet" in dataset_name:
            # masking the predictions corresponding to black pixels around the edges of the image
            black_mask = np.mean(keyframe_image.astype(float), axis=-1) < 10.0
            combined_mask = np.logical_and(black_mask, edge_mask)
            keyframe_prediction[combined_mask] = 0.0

        # masking the depth values larger than max_depth
        keyframe_prediction[keyframe_prediction > max_depth] = 0.0

        keyframe_predictions.append(keyframe_prediction)
        keyframe_images.append(keyframe_image.astype(np.uint8))

    #####
    ##### CALCULATE TSDF VOLUME BOUNDS WITH OR WITHOUT GROUNDTRUTH ANCHORING #####
    #####
    all_groundtruths = None
    if use_groundtruth_to_anchor or save_groundtruth:
        all_groundtruth_filenames = sorted((scene_folder / 'depth').files("*.png"))
        all_groundtruths = []
        for groundtruth_filename in all_groundtruth_filenames:
            groundtruth = cv2.imread(groundtruth_filename, -1).astype(float) / 1000.0
            groundtruth[groundtruth > max_depth] = 0.0
            all_groundtruths.append(groundtruth)

    if use_groundtruth_to_anchor:
        volume_bounds = TSDFFusion.calculate_volume_bounds(all_groundtruths, all_poses, original_K)
    else:
        volume_bounds = TSDFFusion.calculate_volume_bounds(keyframe_predictions, keyframe_poses, scaled_K)
    volume_bounds *= 1.05  # give some margin for the volume bounds accounting for mistakes

    #####
    ##### RUN THE RECONSTRUCTION WITH GROUNDTRUTH DEPTH MAPS #####
    #####
    if save_groundtruth:
        tsdf_volume = TSDFVolume(volume_bounds, voxel_size=voxel_size)

        mesh_name = "{}/reconstruction_voxelsize-{}_maxdepth-{}_anchor-{}_{}_{}_{}".format(reconstruction_folder,
                                                                                          voxel_size,
                                                                                          max_depth,
                                                                                          use_groundtruth_to_anchor,
                                                                                          "GROUNDTRUTH",
                                                                                          dataset_name,
                                                                                          scene_name)
        all_images = []
        for image_filename in all_image_filenames:
            image = load_image(image_filename)
            all_images.append(image)

        TSDFFusion.integrate(tsdf_volume=tsdf_volume,
                             images=all_images,
                             depths=all_groundtruths,
                             poses=all_poses,
                             K=original_K,
                             mesh_name=mesh_name,
                             save_progressive=False)

        tsdf_volume = None
        gc.collect()

    #####
    ##### RUN THE RECONSTRUCTION WITH PREDICTED DEPTH MAPS #####
    #####
    tsdf_volume = TSDFVolume(volume_bounds, voxel_size=voxel_size)

    mesh_name = "{}/reconstruction_voxelsize-{}_maxdepth-{}_anchor-{}_{}_{}_{}".format(reconstruction_folder,
                                                                                       voxel_size,
                                                                                       max_depth,
                                                                                       use_groundtruth_to_anchor,
                                                                                       system_name,
                                                                                       dataset_name,
                                                                                       scene_name)

    TSDFFusion.integrate(tsdf_volume=tsdf_volume,
                         images=keyframe_images,
                         depths=keyframe_predictions,
                         poses=keyframe_poses,
                         K=scaled_K,
                         mesh_name=mesh_name,
                         save_progressive=save_progressive)


if __name__ == "__main__":
    parser = ArgumentParser()

    parser.add_argument('--reconstruction_folder', default="./reconstructions", type=str)
    parser.add_argument('--prediction_folder', default="./predictions", type=str)
    parser.add_argument('--data_folder', default=".", type=str)
    parser.add_argument('--dataset_name', default="hololens-dataset", type=str)
    parser.add_argument('--scene_name', default="000", type=str)
    parser.add_argument('--system_name', default='320_256_3_dvmvs_fusionnet_online', type=str)
    parser.add_argument('--voxel_size', default=0.025, type=float)
    parser.add_argument('--max_depth', default=5.0, type=float)

    parser.add_argument('--use_groundtruth_to_anchor', action='store_true',
                        help="flag for using groundtruth depth maps to calculate the volume bounds and the origin of the volume,"
                             "recommended if groundtruth depth maps are available")
    parser.add_argument('--save_progressive', action='store_true',
                        help="flag for saving the mesh after each fused keyframe to the disk to get the progressive visuals "
                             "similar to the demo video shown in README")
    parser.add_argument('--save_groundtruth', action='store_true',
                        help="flag for saving the complete reconstruction resulting from groundtruth depth maps")

    args = parser.parse_args()

    run(reconstruction_folder=args.reconstruction_folder,
        prediction_folder=args.prediction_folder,
        data_folder=args.data_folder,
        dataset_name=args.dataset_name,
        scene_name=args.scene_name,
        system_name=args.system_name,
        voxel_size=args.voxel_size,
        max_depth=args.max_depth,
        use_groundtruth_to_anchor=args.use_groundtruth_to_anchor,
        save_progressive=args.save_progressive,
        save_groundtruth=args.save_groundtruth)
