-
Notifications
You must be signed in to change notification settings - Fork 97
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Refactor sensor api/camera creation/render mode (#481)
* rename module to sensors * update config * remove _merge_extra_config * refactor code * almost done * fix test and example * add log level to config * test * fix test * format * vehicle panel to panel * cancel extra render frame * fix render bug * fix shadow problem * fix sensor API bug * test pass * panel -> dashboard * use _render_mode to determine * format log message * format * hide debug information * fix bug
- Loading branch information
Showing
59 changed files
with
698 additions
and
662 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
File renamed without changes.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,236 @@ | ||
| import logging | ||
|
|
||
| import numpy as np | ||
|
|
||
| from metadrive.utils.cuda import check_cudart_err | ||
|
|
||
| _cuda_enable = True | ||
| try: | ||
| import cupy as cp | ||
| from OpenGL.GL import GL_TEXTURE_2D # noqa F403 | ||
| from cuda import cudart | ||
| from cuda.cudart import cudaGraphicsRegisterFlags | ||
| from panda3d.core import GraphicsOutput, Texture, GraphicsStateGuardianBase, DisplayRegionDrawCallbackData | ||
| except ImportError: | ||
| _cuda_enable = False | ||
| from panda3d.core import Vec3 | ||
|
|
||
| from metadrive.engine.core.image_buffer import ImageBuffer | ||
|
|
||
|
|
||
| class BaseCamera(ImageBuffer): | ||
| """ | ||
| To enable the image observation, set image_observation to True. | ||
| Every objects share the same camera, to boost the efficiency and save memory. | ||
| Camera configuration is read from the global config automatically. | ||
| """ | ||
| # shape(dim_1, dim_2) | ||
| BUFFER_W = 84 # dim 1 | ||
| BUFFER_H = 84 # dim 2 | ||
| CAM_MASK = None | ||
| display_region_size = [1 / 3, 2 / 3, 0.8, 1.0] | ||
| attached_object = None | ||
|
|
||
| def __init__(self, engine, setup_pbr=False, need_cuda=False): | ||
| self._enable_cuda = need_cuda | ||
| super(BaseCamera, self).__init__( | ||
| self.BUFFER_W, self.BUFFER_H, Vec3(0., 0.8, 1.5), self.BKG_COLOR, setup_pbr=setup_pbr, engine=engine | ||
| ) | ||
|
|
||
| width = self.BUFFER_W | ||
| height = self.BUFFER_H | ||
| if (width > 100 or height > 100) and not self.enable_cuda: | ||
| # Too large height or width will cause corruption in Mac. | ||
| logging.warning( | ||
| "You may using too large buffer! The height is {}, and width is {}. " | ||
| "It may lower the sample efficiency! Considering reduce buffer size or using cuda image by" | ||
| " set [image_on_cuda=True].".format(height, width) | ||
| ) | ||
| self.cuda_graphics_resource = None | ||
| if self.enable_cuda: | ||
| assert _cuda_enable, "Can not enable cuda rendering pipeline" | ||
|
|
||
| # returned tensor property | ||
| self.cuda_dtype = np.uint8 | ||
| self.cuda_shape = (self.BUFFER_W, self.BUFFER_H) | ||
| self.cuda_strides = None | ||
| self.cuda_order = "C" | ||
|
|
||
| self._cuda_buffer = None | ||
|
|
||
| # make texture | ||
| self.cuda_texture = Texture() | ||
| self.buffer.addRenderTexture(self.cuda_texture, GraphicsOutput.RTMBindOrCopy) | ||
|
|
||
| def _callback_func(cbdata: DisplayRegionDrawCallbackData): | ||
| # print("DRAW CALLBACK!!!!!!!!!!!!!!!11") | ||
| cbdata.upcall() | ||
| if not self.registered and self.texture_context_future.done(): | ||
| self.register() | ||
| if self.registered: | ||
| with self as array: | ||
| self.cuda_rendered_result = array | ||
|
|
||
| # Fill the buffer due to multi-thread | ||
| for _ in range(3): | ||
| self.engine.graphicsEngine.renderFrame() | ||
| self.cam.node().getDisplayRegion(0).setDrawCallback(_callback_func) | ||
|
|
||
| self.gsg = GraphicsStateGuardianBase.getDefaultGsg() | ||
| self.texture_context_future = self.cuda_texture.prepare(self.gsg.prepared_objects) | ||
| self.cuda_texture_identifier = None | ||
| self.new_cuda_mem_ptr = None | ||
| self.cuda_rendered_result = None | ||
|
|
||
| @property | ||
| def enable_cuda(self): | ||
| return self is not None and self._enable_cuda | ||
|
|
||
| def get_image(self, base_object): | ||
| """ | ||
| Borrow the camera to get observations | ||
| """ | ||
| self.origin.reparentTo(base_object.origin) | ||
| ret = super(BaseCamera, self).get_image() | ||
| self.track(self.attached_object) | ||
| return ret | ||
|
|
||
| def save_image(self, base_object, name="debug.png"): | ||
| img = self.get_image(base_object) | ||
| img.write(name) | ||
|
|
||
| def get_pixels_array(self, base_object, clip=True) -> np.ndarray: | ||
| self.track(base_object) | ||
| if self.enable_cuda: | ||
| assert self.cuda_rendered_result is not None | ||
| ret = self.cuda_rendered_result[..., :-1][..., ::-1][::-1] | ||
| else: | ||
| ret = self.get_rgb_array_cpu() | ||
| if self.engine.global_config["rgb_to_grayscale"]: | ||
| ret = np.dot(ret[..., :3], [0.299, 0.587, 0.114]) | ||
| if not clip: | ||
| return ret.astype(np.uint8) | ||
| else: | ||
| return ret / 255 | ||
|
|
||
| def destroy(self): | ||
| if self.registered: | ||
| self.unregister() | ||
| ImageBuffer.destroy(self) | ||
|
|
||
| def get_cam(self): | ||
| return self.cam | ||
|
|
||
| def get_lens(self): | ||
| return self.lens | ||
|
|
||
| # # following functions are for onscreen render | ||
| # def add_display_region(self, display_region): | ||
| # super(BaseCamera, self).add_display_region(display_region) | ||
|
|
||
| def remove_display_region(self): | ||
| super(BaseCamera, self).remove_display_region() | ||
|
|
||
| def track(self, base_object): | ||
| if base_object is not None and self is not None: | ||
| self.attached_object = base_object | ||
| self.origin.reparentTo(base_object.origin) | ||
|
|
||
| def __del__(self): | ||
| if self.enable_cuda: | ||
| self.unregister() | ||
| super(BaseCamera, self).__del__() | ||
|
|
||
| """ | ||
| Following functions are cuda support | ||
| """ | ||
|
|
||
| @property | ||
| def cuda_array(self): | ||
| assert self.mapped | ||
| return cp.ndarray( | ||
| shape=(self.cuda_shape[1], self.cuda_shape[0], 4), | ||
| dtype=self.cuda_dtype, | ||
| strides=self.cuda_strides, | ||
| order=self.cuda_order, | ||
| memptr=self._cuda_buffer | ||
| ) | ||
|
|
||
| @property | ||
| def cuda_buffer(self): | ||
| assert self.mapped | ||
| return self._cuda_buffer | ||
|
|
||
| @property | ||
| def graphics_resource(self): | ||
| assert self.registered | ||
| return self.cuda_graphics_resource | ||
|
|
||
| @property | ||
| def registered(self): | ||
| return self.cuda_graphics_resource is not None | ||
|
|
||
| @property | ||
| def mapped(self): | ||
| return self._cuda_buffer is not None | ||
|
|
||
| def __enter__(self): | ||
| return self.map() | ||
|
|
||
| def __exit__(self, exc_type, exc_value, trace): | ||
| self.unmap() | ||
| return False | ||
|
|
||
| def register(self): | ||
| self.cuda_texture_identifier = self.texture_context_future.result().getNativeId() | ||
| assert self.cuda_texture_identifier is not None | ||
| if self.registered: | ||
| return self.cuda_graphics_resource | ||
| self.cuda_graphics_resource = check_cudart_err( | ||
| cudart.cudaGraphicsGLRegisterImage( | ||
| self.cuda_texture_identifier, GL_TEXTURE_2D, cudaGraphicsRegisterFlags.cudaGraphicsRegisterFlagsReadOnly | ||
| ) | ||
| ) | ||
| return self.cuda_graphics_resource | ||
|
|
||
| def unregister(self): | ||
| if self.registered: | ||
| self.unmap() | ||
| self.cuda_graphics_resource = check_cudart_err( | ||
| cudart.cudaGraphicsUnregisterResource(self.cuda_graphics_resource) | ||
| ) | ||
| self.cam.node().getDisplayRegion(0).clearDrawCallback() | ||
|
|
||
| def map(self, stream=0): | ||
| if not self.registered: | ||
| raise RuntimeError("Cannot map an unregistered buffer.") | ||
| if self.mapped: | ||
| return self._cuda_buffer | ||
| check_cudart_err(cudart.cudaGraphicsMapResources(1, self.cuda_graphics_resource, stream)) | ||
| array = check_cudart_err(cudart.cudaGraphicsSubResourceGetMappedArray(self.graphics_resource, 0, 0)) | ||
| channelformat, cudaextent, flag = check_cudart_err(cudart.cudaArrayGetInfo(array)) | ||
|
|
||
| depth = 1 | ||
| byte = 4 # four channel | ||
| if self.new_cuda_mem_ptr is None: | ||
| success, self.new_cuda_mem_ptr = cudart.cudaMalloc(cudaextent.height * cudaextent.width * byte * depth) | ||
| check_cudart_err( | ||
| cudart.cudaMemcpy2DFromArray( | ||
| self.new_cuda_mem_ptr, cudaextent.width * byte * depth, array, 0, 0, cudaextent.width * byte * depth, | ||
| cudaextent.height, cudart.cudaMemcpyKind.cudaMemcpyDeviceToDevice | ||
| ) | ||
| ) | ||
| if self._cuda_buffer is None: | ||
| self._cuda_buffer = cp.cuda.MemoryPointer( | ||
| cp.cuda.UnownedMemory(self.new_cuda_mem_ptr, cudaextent.width * depth * byte * cudaextent.height, self), | ||
| 0 | ||
| ) | ||
| return self.cuda_array | ||
|
|
||
| def unmap(self, stream=None): | ||
| if not self.registered: | ||
| raise RuntimeError("Cannot unmap an unregistered buffer.") | ||
| if not self.mapped: | ||
| return self | ||
| self._cuda_buffer = check_cudart_err(cudart.cudaGraphicsUnmapResources(1, self.cuda_graphics_resource, stream)) | ||
| return self |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.