OpenMotionLab/MotionGPT

1"""Wrapper for offscreen rendering.2 3Author: Matthew Matl4"""5import os6 7from .renderer import Renderer8from .constants import RenderFlags9 10 11class OffscreenRenderer(object):12    """A wrapper for offscreen rendering.13 14    Parameters15    ----------16    viewport_width : int17        The width of the main viewport, in pixels.18    viewport_height : int19        The height of the main viewport, in pixels.20    point_size : float21        The size of screen-space points in pixels.22    """23 24    def __init__(self, viewport_width, viewport_height, point_size=1.0):25        self.viewport_width = viewport_width26        self.viewport_height = viewport_height27        self.point_size = point_size28 29        self._platform = None30        self._renderer = None31        self._create()32 33    @property34    def viewport_width(self):35        """int : The width of the main viewport, in pixels.36        """37        return self._viewport_width38 39    @viewport_width.setter40    def viewport_width(self, value):41        self._viewport_width = int(value)42 43    @property44    def viewport_height(self):45        """int : The height of the main viewport, in pixels.46        """47        return self._viewport_height48 49    @viewport_height.setter50    def viewport_height(self, value):51        self._viewport_height = int(value)52 53    @property54    def point_size(self):55        """float : The pixel size of points in point clouds.56        """57        return self._point_size58 59    @point_size.setter60    def point_size(self, value):61        self._point_size = float(value)62 63    def render(self, scene, flags=RenderFlags.NONE, seg_node_map=None):64        """Render a scene with the given set of flags.65 66        Parameters67        ----------68        scene : :class:`Scene`69            A scene to render.70        flags : int71            A bitwise or of one or more flags from :class:`.RenderFlags`.72        seg_node_map : dict73            A map from :class:`.Node` objects to (3,) colors for each.74            If specified along with flags set to :attr:`.RenderFlags.SEG`,75            the color image will be a segmentation image.76 77        Returns78        -------79        color_im : (h, w, 3) uint8 or (h, w, 4) uint880            The color buffer in RGB format, or in RGBA format if81            :attr:`.RenderFlags.RGBA` is set.82            Not returned if flags includes :attr:`.RenderFlags.DEPTH_ONLY`.83        depth_im : (h, w) float3284            The depth buffer in linear units.85        """86        self._platform.make_current()87        # If platform does not support dynamically-resizing framebuffers,88        # destroy it and restart it89        if (self._platform.viewport_height != self.viewport_height or90                self._platform.viewport_width != self.viewport_width):91            if not self._platform.supports_framebuffers():92                self.delete()93                self._create()94 95        self._platform.make_current()96        self._renderer.viewport_width = self.viewport_width97        self._renderer.viewport_height = self.viewport_height98        self._renderer.point_size = self.point_size99 100        if self._platform.supports_framebuffers():101            flags |= RenderFlags.OFFSCREEN102            retval = self._renderer.render(scene, flags, seg_node_map)103        else:104            self._renderer.render(scene, flags, seg_node_map)105            depth = self._renderer.read_depth_buf()106            if flags & RenderFlags.DEPTH_ONLY:107                retval = depth108            else:109                color = self._renderer.read_color_buf()110                retval = color, depth111 112        # Make the platform not current113        self._platform.make_uncurrent()114        return retval115 116    def delete(self):117        """Free all OpenGL resources.118        """119        self._platform.make_current()120        self._renderer.delete()121        self._platform.delete_context()122        del self._renderer123        del self._platform124        self._renderer = None125        self._platform = None126        import gc127        gc.collect()128 129    def _create(self):130        if 'PYOPENGL_PLATFORM' not in os.environ:131            from pyrender.platforms.pyglet_platform import PygletPlatform132            self._platform = PygletPlatform(self.viewport_width,133                                            self.viewport_height)134        elif os.environ['PYOPENGL_PLATFORM'] == 'egl':135            from pyrender.platforms import egl136            device_id = int(os.environ.get('EGL_DEVICE_ID', '0'))137            egl_device = egl.get_device_by_index(device_id)138            self._platform = egl.EGLPlatform(self.viewport_width,139                                             self.viewport_height,140                                             device=egl_device)141        elif os.environ['PYOPENGL_PLATFORM'] == 'osmesa':142            from pyrender.platforms.osmesa import OSMesaPlatform143            self._platform = OSMesaPlatform(self.viewport_width,144                                            self.viewport_height)145        else:146            raise ValueError('Unsupported PyOpenGL platform: {}'.format(147                os.environ['PYOPENGL_PLATFORM']148            ))149        self._platform.init_context()150        self._platform.make_current()151        self._renderer = Renderer(self.viewport_width, self.viewport_height)152 153    def __del__(self):154        try:155            self.delete()156        except Exception:157            pass158 159 160__all__ = ['OffscreenRenderer']161