def read_pixels()

in mujoco_py/mjrendercontext.pyx [0:0]

• 30 lines of code
• 5 McCabe index (conditional complexity)

    def read_pixels(self, width, height, depth=True, segmentation=False):
        cdef mjrRect rect
        rect.left = 0
        rect.bottom = 0
        rect.width = width
        rect.height = height

        rgb_arr = np.zeros(3 * rect.width * rect.height, dtype=np.uint8)
        depth_arr = np.zeros(rect.width * rect.height, dtype=np.float32)
        
        cdef unsigned char[::view.contiguous] rgb_view = rgb_arr
        cdef float[::view.contiguous] depth_view = depth_arr
        mjr_readPixels(&rgb_view[0], &depth_view[0], rect, &self._con)
        rgb_img = rgb_arr.reshape(rect.height, rect.width, 3)
        cdef np.ndarray[np.npy_uint32, ndim=2] seg_img
        cdef np.ndarray[np.npy_int32, ndim=2] seg_ids

        ret_img = rgb_img
        if segmentation:
            seg_img = (rgb_img[:, :, 0] + rgb_img[:, :, 1] * (2**8) + rgb_img[:, :, 2] * (2 ** 16))
            seg_img[seg_img >= (self._scn.ngeom + 1)] = 0
            seg_ids = np.full((self._scn.ngeom + 1, 2), fill_value=-1, dtype=np.int32)

            for i in range(self._scn.ngeom):
                geom = self._scn.geoms[i]
                if geom.segid != -1:
                    seg_ids[geom.segid + 1, 0] = geom.objtype
                    seg_ids[geom.segid + 1, 1] = geom.objid
            ret_img = seg_ids[seg_img]

        if depth:
            depth_img = depth_arr.reshape(rect.height, rect.width)
            return (ret_img, depth_img)
        else:
            return ret_img