def forward(self, depth, pose=torch.eye(4).float()): self.device = depth.device # Works for FOV = 90 degrees # Should be adjusted, if FOV changed self.fx = float(depth.size(1)) / 2.0 self.fy = float(depth.size(0)) / 2.0 self.cx = int(self.fx) - 1 self.cy = int(self.fy) - 1 pose = pose.to(self.device) local_3d_pcl = depth2local3d(depth, self.fx, self.fy, self.cx, self.cy) idxs = (torch.abs(local_3d_pcl[:, 2]) < self.far_th) * ( torch.abs(local_3d_pcl[:, 2]) >= self.near_th ) survived_points = local_3d_pcl[idxs] if len(survived_points) < 20: map_size_in_cells = ( get_map_size_in_cells(self.map_size_meters, self.map_cell_size) - 1 ) init_map = torch.zeros( (map_size_in_cells, map_size_in_cells), device=self.device ) return init_map global_3d_pcl = reproject_local_to_global(survived_points, pose)[:, :3] # Because originally y looks down and from agent camera height global_3d_pcl[:, 1] = -global_3d_pcl[:, 1] + self.camera_height idxs = (global_3d_pcl[:, 1] > self.h_min_th) * ( global_3d_pcl[:, 1] < self.h_max_th ) global_3d_pcl = global_3d_pcl[idxs] obstacle_map = pcl_to_obstacles( global_3d_pcl, self.map_size_meters, self.map_cell_size ) return obstacle_map
def pcl_to_obstacles(pts3d, map_size=40, cell_size=0.2, min_pts=10): r"""Counts number of 3d points in 2d map cell. Height is sum-pooled. """ device = pts3d.device map_size_in_cells = get_map_size_in_cells(map_size, cell_size) - 1 init_map = torch.zeros((map_size_in_cells, map_size_in_cells), device=device) if len(pts3d) <= 1: return init_map num_pts, dim = pts3d.size() pts2d = torch.cat([pts3d[:, 2:3], pts3d[:, 0:1]], dim=1) data_idxs = torch.round( project2d_pcl_into_worldmap(pts2d, map_size, cell_size)) if len(data_idxs) > min_pts: u, counts = np.unique(data_idxs.detach().cpu().numpy(), axis=0, return_counts=True) init_map[u[:, 0], u[:, 1]] = torch.from_numpy(counts).to(dtype=torch.float32, device=device) return init_map