diff --git a/active_grasp/nbv.py b/active_grasp/nbv.py index 9e1f11a..48da9cc 100644 --- a/active_grasp/nbv.py +++ b/active_grasp/nbv.py @@ -1,10 +1,56 @@ import itertools +from numba import jit import numpy as np import rospy -from .timer import Timer from .policy import MultiViewPolicy +from .timer import Timer + + +@jit(nopython=True) +def get_voxel_at(voxel_size, p): + index = (p / voxel_size).astype(np.int64) + return index if (index >= 0).all() and (index < 40).all() else None + + +@jit(nopython=True) +def raycast( + voxel_size, + tsdf_grid, + ori, + pos, + fx, + fy, + cx, + cy, + u_min, + u_max, + v_min, + v_max, + t_min, + t_max, + t_step, +): + voxel_indices = [] + for u in range(u_min, u_max): + for v in range(v_min, v_max): + direction = np.asarray([(u - cx) / fx, (v - cy) / fy, 1.0]) + direction = ori @ (direction / np.linalg.norm(direction)) + t, tsdf_prev = t_min, -1.0 + while t < t_max: + p = pos + t * direction + t += t_step + index = get_voxel_at(voxel_size, p) + if index is not None: + i, j, k = index + tsdf = tsdf_grid[i, j, k] + if tsdf * tsdf_prev < 0 and tsdf_prev > -1: # crossed a surface + break + voxel_indices.append(index) + tsdf_prev = tsdf + return voxel_indices + class NextBestView(MultiViewPolicy): def __init__(self): @@ -61,6 +107,7 @@ class NextBestView(MultiViewPolicy): tsdf_grid, voxel_size = self.tsdf.get_grid(), self.tsdf.voxel_size tsdf_grid = -1.0 + 2.0 * tsdf_grid # Open3D maps tsdf to [0,1] + # Downsample the sensor resolution fx = self.intrinsic.fx / downsample fy = self.intrinsic.fy / downsample cx = self.intrinsic.cx / downsample @@ -78,32 +125,27 @@ class NextBestView(MultiViewPolicy): t_max = corners[2].max() # TODO This bound might be a bit too short t_step = np.sqrt(3) * voxel_size # TODO replace with line rasterization - view = self.T_task_base * view # We'll work in the task frame from now on - origin = view.translation + # Cast rays from the camera view (we'll work in the task frame from now on) + view = self.T_task_base * view + ori, pos = view.rotation.as_matrix(), view.translation - def get_voxel_at(p): - index = (p / voxel_size).astype(int) - return index if (index >= 0).all() and (index < 40).all() else None - - voxel_indices = [] - for u in range(u_min, u_max): - for v in range(v_min, v_max): - direction = np.r_[(u - cx) / fx, (v - cy) / fy, 1.0] - direction = view.rotation.apply(direction / np.linalg.norm(direction)) - # self.vis.rays(self.task_frame, origin, [direction]) - t, tsdf_prev = t_min, -1.0 - while t < t_max: - p = origin + t * direction - t += t_step - # self.vis.point(self.task_frame, p) - index = get_voxel_at(p) - if index is not None: - i, j, k = index - tsdf = tsdf_grid[i, j, k] - if tsdf * tsdf_prev < 0 and tsdf_prev > -1: # Crossed a surface - break - voxel_indices.append(index) - tsdf_prev = tsdf + voxel_indices = raycast( + voxel_size, + tsdf_grid, + ori, + pos, + fx, + fy, + cx, + cy, + u_min, + u_max, + v_min, + v_max, + t_min, + t_max, + t_step, + ) # Count rear side voxels i, j, k = np.unique(voxel_indices, axis=0).T diff --git a/requirements.txt b/requirements.txt index 391ca2f..a5f76bc 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,2 @@ +numba scikit-image