|
| 1 | +""" |
| 2 | +CUDA backend adapter for OneMap. Wraps cuonemap.OneMap to present the same interface as |
| 3 | +the Python OneMap class in feature_map.py, so the two are interchangeable. |
| 4 | +
|
| 5 | +See https://github.com/finnBsch/CuOneMap for the backend itself, including the known |
| 6 | +differences from the PyTorch implementation. |
| 7 | +""" |
| 8 | +import numpy as np |
| 9 | +import torch |
| 10 | + |
| 11 | +import cuonemap as _onemap_cuda |
| 12 | + |
| 13 | + |
| 14 | +class CUDAOneMap: |
| 15 | + """Drop-in replacement for OneMap using the CUDA backend. |
| 16 | +
|
| 17 | + Accepts the same constructor arguments and exposes the same properties/methods |
| 18 | + as the Python OneMap class, translating between the Python and CUDA interfaces |
| 19 | + internally. |
| 20 | + """ |
| 21 | + |
| 22 | + def __init__(self, |
| 23 | + feature_dim, |
| 24 | + config, |
| 25 | + dense_projection=None, |
| 26 | + fusion_type=None, |
| 27 | + map_device="cuda"): |
| 28 | + assert feature_dim == 768, ( |
| 29 | + f"CUDA OneMap backend requires feature_dim=768 (CLIP), got {feature_dim}" |
| 30 | + ) |
| 31 | + cuda_conf = _onemap_cuda.MappingConf() |
| 32 | + cuda_conf.n_cells = config.n_points |
| 33 | + cuda_conf.size = float(config.size) |
| 34 | + cuda_conf.agent_radius = config.agent_radius |
| 35 | + cuda_conf.blur_kernel_size = config.blur_kernel_size |
| 36 | + cuda_conf.obstacle_map_threshold = config.obstacle_map_threshold |
| 37 | + cuda_conf.fully_explored_threshold = config.fully_explored_threshold |
| 38 | + cuda_conf.checked_map_threshold = config.checked_map_threshold |
| 39 | + cuda_conf.depth_factor = config.depth_factor |
| 40 | + cuda_conf.gradient_factor = config.gradient_factor |
| 41 | + cuda_conf.optimal_object_distance = config.optimal_object_distance |
| 42 | + cuda_conf.optimal_object_factor = config.optimal_object_factor |
| 43 | + cuda_conf.obstacle_min = config.obstacle_min |
| 44 | + cuda_conf.obstacle_max = config.obstacle_max |
| 45 | + cuda_conf.filter_stairs = config.filter_stairs |
| 46 | + cuda_conf.floor_level = config.floor_level |
| 47 | + cuda_conf.floor_threshold = config.floor_threshold |
| 48 | + cuda_conf.use_cpu_storage = (map_device == "cpu") |
| 49 | + |
| 50 | + self._impl = _onemap_cuda.OneMap(cuda_conf) |
| 51 | + self.map_device = map_device |
| 52 | + self.feature_dim = feature_dim |
| 53 | + self.n_cells = config.n_points |
| 54 | + self.size = float(config.size) |
| 55 | + self.cell_size = self.size / self.n_cells |
| 56 | + center = self._impl.map_center_cells |
| 57 | + self._map_center_cells = torch.tensor( |
| 58 | + [center[0], center[1]], dtype=torch.int32 |
| 59 | + ).to("cuda") |
| 60 | + |
| 61 | + @property |
| 62 | + def map_center_cells(self): |
| 63 | + return self._map_center_cells |
| 64 | + |
| 65 | + @property |
| 66 | + def camera_initialized(self): |
| 67 | + return self._impl.camera_initialized |
| 68 | + |
| 69 | + @property |
| 70 | + def fx(self): |
| 71 | + return self._impl.fx |
| 72 | + |
| 73 | + @property |
| 74 | + def fy(self): |
| 75 | + return self._impl.fy |
| 76 | + |
| 77 | + @property |
| 78 | + def cx(self): |
| 79 | + return self._impl.cx |
| 80 | + |
| 81 | + @property |
| 82 | + def cy(self): |
| 83 | + return self._impl.cy |
| 84 | + |
| 85 | + # --- Map properties --- |
| 86 | + |
| 87 | + @property |
| 88 | + def feature_map(self): |
| 89 | + return self._impl.feature_map |
| 90 | + |
| 91 | + @property |
| 92 | + def confidence_map(self): |
| 93 | + return self._impl.confidence_map |
| 94 | + |
| 95 | + @property |
| 96 | + def checked_conf_map(self): |
| 97 | + return self._impl.checked_conf_map |
| 98 | + |
| 99 | + @property |
| 100 | + def obstacle_map(self): |
| 101 | + return self._impl.obstacle_map |
| 102 | + |
| 103 | + @property |
| 104 | + def updated_mask(self): |
| 105 | + return self._impl.updated_mask |
| 106 | + |
| 107 | + @property |
| 108 | + def navigable_map(self): |
| 109 | + t = self._impl.navigable_map |
| 110 | + if isinstance(t, np.ndarray): |
| 111 | + return t |
| 112 | + return t.cpu().numpy() if hasattr(t, 'cpu') else np.asarray(t) |
| 113 | + |
| 114 | + @property |
| 115 | + def fully_explored_map(self): |
| 116 | + t = self._impl.fully_explored_map |
| 117 | + if isinstance(t, np.ndarray): |
| 118 | + return t |
| 119 | + return t.cpu().numpy() if hasattr(t, 'cpu') else np.asarray(t) |
| 120 | + |
| 121 | + @property |
| 122 | + def checked_map(self): |
| 123 | + t = self._impl.checked_map |
| 124 | + if isinstance(t, np.ndarray): |
| 125 | + return t |
| 126 | + return t.cpu().numpy() if hasattr(t, 'cpu') else np.asarray(t) |
| 127 | + |
| 128 | + @property |
| 129 | + def occluded_map(self): |
| 130 | + t = self._impl.occluded_map |
| 131 | + if isinstance(t, np.ndarray): |
| 132 | + return t |
| 133 | + return t.cpu().numpy() if hasattr(t, 'cpu') else np.asarray(t) |
| 134 | + |
| 135 | + # --- Methods --- |
| 136 | + |
| 137 | + def set_camera_matrix(self, camera_matrix): |
| 138 | + self._impl.set_camera_matrix(camera_matrix) |
| 139 | + |
| 140 | + def update(self, values, depth, tf_camera_to_episodic, artifical_obstacles=None): |
| 141 | + # values: the Python implementation is given (F, H, W); the backend wants (H, W, F) |
| 142 | + if len(values.shape) == 3 and values.shape[0] == self.feature_dim: |
| 143 | + values = values.permute(1, 2, 0).contiguous() |
| 144 | + if not values.is_cuda: |
| 145 | + values = values.to("cuda") |
| 146 | + |
| 147 | + # depth: numpy or CPU tensor -> float32 CUDA tensor |
| 148 | + if isinstance(depth, np.ndarray): |
| 149 | + depth = torch.from_numpy(depth.astype(np.float32)).to("cuda") |
| 150 | + elif isinstance(depth, torch.Tensor): |
| 151 | + if not depth.is_cuda: |
| 152 | + depth = depth.to(torch.float32).to("cuda") |
| 153 | + else: |
| 154 | + depth = torch.as_tensor(depth, dtype=torch.float32).to("cuda") |
| 155 | + |
| 156 | + # tf stays on the host: the binding takes a CPU array |
| 157 | + if isinstance(tf_camera_to_episodic, torch.Tensor): |
| 158 | + tf_camera_to_episodic = tf_camera_to_episodic.cpu().numpy() |
| 159 | + tf_camera_to_episodic = tf_camera_to_episodic.astype(np.float32) |
| 160 | + |
| 161 | + obstacles = [] |
| 162 | + if artifical_obstacles: |
| 163 | + obstacles = [(float(o[0]), float(o[1])) for o in artifical_obstacles] |
| 164 | + |
| 165 | + self._impl.update(values, depth, tf_camera_to_episodic, obstacles) |
| 166 | + |
| 167 | + def reset(self): |
| 168 | + self._impl.reset() |
| 169 | + |
| 170 | + def reset_updated_mask(self): |
| 171 | + self._impl.reset_updated_mask() |
| 172 | + |
| 173 | + def reset_checked_map(self): |
| 174 | + self._impl.reset_checked_map() |
| 175 | + |
| 176 | + def metric_to_px(self, x, y): |
| 177 | + return self._impl.metric_to_px(x, y) |
| 178 | + |
| 179 | + def px_to_metric(self, px, py): |
| 180 | + return self._impl.px_to_metric(px, py) |
0 commit comments