def get_observation(self): tgt = self._sensor_object.render_target if self._spec.gpu2gpu_transfer: with torch.cuda.device(self._buffer.device): if self._spec.sensor_type == hsim.SensorType.SEMANTIC: tgt.read_frame_object_id_gpu(self._buffer.data_ptr()) elif self._spec.sensor_type == hsim.SensorType.DEPTH: tgt.read_frame_depth_gpu(self._buffer.data_ptr()) else: tgt.read_frame_rgba_gpu(self._buffer.data_ptr()) return self._buffer.flip(0).clone() else: size = self._sensor_object.framebuffer_size if self._spec.sensor_type == hsim.SensorType.SEMANTIC: tgt.read_frame_object_id( mn.MutableImageView2D(mn.PixelFormat.R32UI, size, self._buffer)) elif self._spec.sensor_type == hsim.SensorType.DEPTH: tgt.read_frame_depth( mn.MutableImageView2D(mn.PixelFormat.R32F, size, self._buffer)) else: tgt.read_frame_rgba( mn.MutableImageView2D( mn.PixelFormat.RGBA8UNORM, size, self._buffer.reshape(self._spec.resolution[0], -1), )) return np.flip(self._buffer, axis=0).copy()
def get_observation(self) -> Union[ndarray, "Tensor"]: tgt = self._sensor_object.render_target if self._spec.gpu2gpu_transfer: with torch.cuda.device( self._buffer.device): # type: ignore[attr-defined] if self._spec.sensor_type == SensorType.SEMANTIC: tgt.read_frame_object_id_gpu( self._buffer.data_ptr()) # type: ignore[attr-defined] elif self._spec.sensor_type == SensorType.DEPTH: tgt.read_frame_depth_gpu( self._buffer.data_ptr()) # type: ignore[attr-defined] else: tgt.read_frame_rgba_gpu( self._buffer.data_ptr()) # type: ignore[attr-defined] obs = self._buffer.flip(0) else: size = self._sensor_object.framebuffer_size if self._spec.sensor_type == SensorType.SEMANTIC: tgt.read_frame_object_id( mn.MutableImageView2D(mn.PixelFormat.R32UI, size, self._buffer)) elif self._spec.sensor_type == SensorType.DEPTH: tgt.read_frame_depth( mn.MutableImageView2D(mn.PixelFormat.R32F, size, self._buffer)) else: tgt.read_frame_rgba( mn.MutableImageView2D( mn.PixelFormat.RGBA8_UNORM, size, self._buffer.reshape(self._spec.resolution[0], -1), )) obs = np.flip(self._buffer, axis=0) return self._noise_model(obs)
def __init__(self, sim: Simulator, agent: Agent, sensor_id: str) -> None: self._sim = sim self._agent = agent # sensor is an attached object to the scene node # store such "attached object" in _sensor_object self._sensor_object = self._agent._sensors[sensor_id] self._spec = self._sensor_object.specification() if self._sim.renderer is not None: self._sim.renderer.bind_render_target(self._sensor_object) if self._spec.gpu2gpu_transfer: assert cuda_enabled, "Must build habitat sim with cuda for gpu2gpu-transfer" assert _HAS_TORCH device = torch.device( "cuda", self._sim.gpu_device) # type: ignore[attr-defined] torch.cuda.set_device(device) resolution = self._spec.resolution if self._spec.sensor_type == SensorType.SEMANTIC: self._buffer: Union[np.ndarray, "Tensor"] = torch.empty(resolution[0], resolution[1], dtype=torch.int32, device=device) elif self._spec.sensor_type == SensorType.DEPTH: self._buffer = torch.empty(resolution[0], resolution[1], dtype=torch.float32, device=device) else: self._buffer = torch.empty(resolution[0], resolution[1], 4, dtype=torch.uint8, device=device) else: size = self._sensor_object.framebuffer_size if self._spec.sensor_type == SensorType.SEMANTIC: self._buffer = np.empty( (self._spec.resolution[0], self._spec.resolution[1]), dtype=np.uint32, ) self.view = mn.MutableImageView2D(mn.PixelFormat.R32UI, size, self._buffer) elif self._spec.sensor_type == SensorType.DEPTH: self._buffer = np.empty( (self._spec.resolution[0], self._spec.resolution[1]), dtype=np.float32, ) self.view = mn.MutableImageView2D(mn.PixelFormat.R32F, size, self._buffer) else: self._buffer = np.empty( ( self._spec.resolution[0], self._spec.resolution[1], self._spec.channels, ), dtype=np.uint8, ) self.view = mn.MutableImageView2D( mn.PixelFormat.RGBA8_UNORM, size, self._buffer.reshape(self._spec.resolution[0], -1), ) noise_model_kwargs = self._spec.noise_model_kwargs self._noise_model = make_sensor_noise_model( self._spec.noise_model, { "gpu_device_id": self._sim.gpu_device, **noise_model_kwargs }, ) assert self._noise_model.is_valid_sensor_type( self._spec.sensor_type ), "Noise model '{}' is not valid for sensor '{}'".format( self._spec.noise_model, self._spec.uuid)