|
| 1 | +""" |
| 2 | +opencv renderer class. |
| 3 | +""" |
| 4 | +import platform |
| 5 | + |
| 6 | +import cv2 |
| 7 | +import numpy as np |
| 8 | + |
| 9 | + |
| 10 | +class OpenCVViewer: |
| 11 | + def __init__(self, sim): |
| 12 | + # TODO: update this appropriately - need to get screen dimensions |
| 13 | + self.width = 1280 |
| 14 | + self.height = 800 |
| 15 | + |
| 16 | + self.sim = sim |
| 17 | + self.set_camera(camera_id=0) |
| 18 | + self._window_name = "offscreen render" |
| 19 | + self._has_window = False |
| 20 | + self.keypress_callback = None |
| 21 | + |
| 22 | + def set_camera(self, camera_id=None, camera_name=None, width=None, height=None): |
| 23 | + """ |
| 24 | + Set the camera view to the specified camera ID. |
| 25 | +
|
| 26 | + Args: |
| 27 | + camera_id (int or list): id(s) of the camera to set the current viewer to |
| 28 | + camera_name (str or list or None): name(s) of the camera to set the current viewer to |
| 29 | + """ |
| 30 | + |
| 31 | + # enforce exactly one arg |
| 32 | + assert (camera_id is not None) or (camera_name is not None) |
| 33 | + assert (camera_id is None) or (camera_name is None) |
| 34 | + |
| 35 | + # set width and height |
| 36 | + if width is not None: |
| 37 | + self.width = width |
| 38 | + if height is not None: |
| 39 | + self.height = height |
| 40 | + |
| 41 | + if camera_id is not None: |
| 42 | + if isinstance(camera_id, int): |
| 43 | + camera_id = [camera_id] |
| 44 | + self.camera_names = [self.sim.model.camera_id2name(cam_id) for cam_id in camera_id] |
| 45 | + else: |
| 46 | + if isinstance(camera_name, str): |
| 47 | + camera_name = [camera_name] |
| 48 | + self.camera_names = list(camera_name) |
| 49 | + |
| 50 | + def render(self): |
| 51 | + # get frame with offscreen renderer (assumes that the renderer already exists) |
| 52 | + im = [ |
| 53 | + self.sim.render(camera_name=cam_name, height=self.height, width=self.width)[..., ::-1] |
| 54 | + for cam_name in self.camera_names |
| 55 | + ] |
| 56 | + im = np.concatenate(im, axis=1) # concatenate horizontally |
| 57 | + |
| 58 | + # write frame to window |
| 59 | + im = np.flip(im, axis=0) |
| 60 | + cv2.imshow(self._window_name, im) |
| 61 | + if (platform.system() != "Darwin") and (not self._has_window): |
| 62 | + # move window to top left of screen, and ensure we only move window on creation |
| 63 | + cv2.moveWindow(self._window_name, 0, 0) |
| 64 | + key = cv2.waitKey(1) |
| 65 | + if self.keypress_callback is not None: |
| 66 | + self.keypress_callback(key) |
| 67 | + self._has_window = True |
| 68 | + |
| 69 | + def add_keypress_callback(self, keypress_callback): |
| 70 | + self.keypress_callback = keypress_callback |
| 71 | + |
| 72 | + def close_window(self): |
| 73 | + """ |
| 74 | + Helper method to close the active window. |
| 75 | + """ |
| 76 | + if self._has_window: |
| 77 | + cv2.destroyWindow(self._window_name) |
| 78 | + cv2.waitKey(1) |
| 79 | + self._has_window = False |
| 80 | + |
| 81 | + def close(self): |
| 82 | + """ |
| 83 | + Any cleanup to close renderer. |
| 84 | + """ |
| 85 | + |
| 86 | + # NOTE: assume that @sim will get cleaned up outside the renderer - just delete the reference |
| 87 | + self.sim = None |
| 88 | + |
| 89 | + # close window |
| 90 | + self.close_window() |
0 commit comments