facebook/StyleNeRF
34
1# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.2#3# NVIDIA CORPORATION and its licensors retain all intellectual property4# and proprietary rights in and to this software, related documentation5# and any modifications thereto. Any use, reproduction, disclosure or6# distribution of this software and related documentation without an express7# license agreement from NVIDIA CORPORATION is strictly prohibited.8 9import os10import functools11import contextlib12import numpy as np13import OpenGL.GL as gl14import OpenGL.GL.ARB.texture_float15import dnnlib16 17#----------------------------------------------------------------------------18 19def init_egl():20 assert os.environ['PYOPENGL_PLATFORM'] == 'egl' # Must be set before importing OpenGL.21 import OpenGL.EGL as egl22 import ctypes23 24 # Initialize EGL.25 display = egl.eglGetDisplay(egl.EGL_DEFAULT_DISPLAY)26 assert display != egl.EGL_NO_DISPLAY27 major = ctypes.c_int32()28 minor = ctypes.c_int32()29 ok = egl.eglInitialize(display, major, minor)30 assert ok31 assert major.value * 10 + minor.value >= 1432 33 # Choose config.34 config_attribs = [35 egl.EGL_RENDERABLE_TYPE, egl.EGL_OPENGL_BIT,36 egl.EGL_SURFACE_TYPE, egl.EGL_PBUFFER_BIT,37 egl.EGL_NONE38 ]39 configs = (ctypes.c_int32 * 1)()40 num_configs = ctypes.c_int32()41 ok = egl.eglChooseConfig(display, config_attribs, configs, 1, num_configs)42 assert ok43 assert num_configs.value == 144 config = configs[0]45 46 # Create dummy pbuffer surface.47 surface_attribs = [48 egl.EGL_WIDTH, 1,49 egl.EGL_HEIGHT, 1,50 egl.EGL_NONE51 ]52 surface = egl.eglCreatePbufferSurface(display, config, surface_attribs)53 assert surface != egl.EGL_NO_SURFACE54 55 # Setup GL context.56 ok = egl.eglBindAPI(egl.EGL_OPENGL_API)57 assert ok58 context = egl.eglCreateContext(display, config, egl.EGL_NO_CONTEXT, None)59 assert context != egl.EGL_NO_CONTEXT60 ok = egl.eglMakeCurrent(display, surface, surface, context)61 assert ok62 63#----------------------------------------------------------------------------64 65_texture_formats = {66 ('uint8', 1): dnnlib.EasyDict(type=gl.GL_UNSIGNED_BYTE, format=gl.GL_LUMINANCE, internalformat=gl.GL_LUMINANCE8),67 ('uint8', 2): dnnlib.EasyDict(type=gl.GL_UNSIGNED_BYTE, format=gl.GL_LUMINANCE_ALPHA, internalformat=gl.GL_LUMINANCE8_ALPHA8),68 ('uint8', 3): dnnlib.EasyDict(type=gl.GL_UNSIGNED_BYTE, format=gl.GL_RGB, internalformat=gl.GL_RGB8),69 ('uint8', 4): dnnlib.EasyDict(type=gl.GL_UNSIGNED_BYTE, format=gl.GL_RGBA, internalformat=gl.GL_RGBA8),70 ('float32', 1): dnnlib.EasyDict(type=gl.GL_FLOAT, format=gl.GL_LUMINANCE, internalformat=OpenGL.GL.ARB.texture_float.GL_LUMINANCE32F_ARB),71 ('float32', 2): dnnlib.EasyDict(type=gl.GL_FLOAT, format=gl.GL_LUMINANCE_ALPHA, internalformat=OpenGL.GL.ARB.texture_float.GL_LUMINANCE_ALPHA32F_ARB),72 ('float32', 3): dnnlib.EasyDict(type=gl.GL_FLOAT, format=gl.GL_RGB, internalformat=gl.GL_RGB32F),73 ('float32', 4): dnnlib.EasyDict(type=gl.GL_FLOAT, format=gl.GL_RGBA, internalformat=gl.GL_RGBA32F),74}75 76def get_texture_format(dtype, channels):77 return _texture_formats[(np.dtype(dtype).name, int(channels))]78 79#----------------------------------------------------------------------------80 81def prepare_texture_data(image):82 image = np.asarray(image)83 if image.ndim == 2:84 image = image[:, :, np.newaxis]85 if image.dtype.name == 'float64':86 image = image.astype('float32')87 return image88 89#----------------------------------------------------------------------------90 91def draw_pixels(image, *, pos=0, zoom=1, align=0, rint=True):92 pos = np.broadcast_to(np.asarray(pos, dtype='float32'), [2])93 zoom = np.broadcast_to(np.asarray(zoom, dtype='float32'), [2])94 align = np.broadcast_to(np.asarray(align, dtype='float32'), [2])95 image = prepare_texture_data(image)96 height, width, channels = image.shape97 size = zoom * [width, height]98 pos = pos - size * align99 if rint:100 pos = np.rint(pos)101 fmt = get_texture_format(image.dtype, channels)102 103 gl.glPushAttrib(gl.GL_CURRENT_BIT | gl.GL_PIXEL_MODE_BIT)104 gl.glPushClientAttrib(gl.GL_CLIENT_PIXEL_STORE_BIT)105 gl.glRasterPos2f(pos[0], pos[1])106 gl.glPixelZoom(zoom[0], -zoom[1])107 gl.glPixelStorei(gl.GL_UNPACK_ALIGNMENT, 1)108 gl.glDrawPixels(width, height, fmt.format, fmt.type, image)109 gl.glPopClientAttrib()110 gl.glPopAttrib()111 112#----------------------------------------------------------------------------113 114def read_pixels(width, height, *, pos=0, dtype='uint8', channels=3):115 pos = np.broadcast_to(np.asarray(pos, dtype='float32'), [2])116 dtype = np.dtype(dtype)117 fmt = get_texture_format(dtype, channels)118 image = np.empty([height, width, channels], dtype=dtype)119 120 gl.glPushClientAttrib(gl.GL_CLIENT_PIXEL_STORE_BIT)121 gl.glPixelStorei(gl.GL_PACK_ALIGNMENT, 1)122 gl.glReadPixels(int(np.round(pos[0])), int(np.round(pos[1])), width, height, fmt.format, fmt.type, image)123 gl.glPopClientAttrib()124 return np.flipud(image)125 126#----------------------------------------------------------------------------127 128class Texture:129 def __init__(self, *, image=None, width=None, height=None, channels=None, dtype=None, bilinear=True, mipmap=True):130 self.gl_id = None131 self.bilinear = bilinear132 self.mipmap = mipmap133 134 # Determine size and dtype.135 if image is not None:136 image = prepare_texture_data(image)137 self.height, self.width, self.channels = image.shape138 self.dtype = image.dtype139 else:140 assert width is not None and height is not None141 self.width = width142 self.height = height143 self.channels = channels if channels is not None else 3144 self.dtype = np.dtype(dtype) if dtype is not None else np.uint8145 146 # Validate size and dtype.147 assert isinstance(self.width, int) and self.width >= 0148 assert isinstance(self.height, int) and self.height >= 0149 assert isinstance(self.channels, int) and self.channels >= 1150 assert self.is_compatible(width=width, height=height, channels=channels, dtype=dtype)151 152 # Create texture object.153 self.gl_id = gl.glGenTextures(1)154 with self.bind():155 gl.glTexParameterf(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_S, gl.GL_CLAMP_TO_EDGE)156 gl.glTexParameterf(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_T, gl.GL_CLAMP_TO_EDGE)157 gl.glTexParameterf(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MAG_FILTER, gl.GL_LINEAR if self.bilinear else gl.GL_NEAREST)158 gl.glTexParameterf(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MIN_FILTER, gl.GL_LINEAR_MIPMAP_LINEAR if self.mipmap else gl.GL_NEAREST)159 self.update(image)160 161 def delete(self):162 if self.gl_id is not None:163 gl.glDeleteTextures([self.gl_id])164 self.gl_id = None165 166 def __del__(self):167 try:168 self.delete()169 except:170 pass171 172 @contextlib.contextmanager173 def bind(self):174 prev_id = gl.glGetInteger(gl.GL_TEXTURE_BINDING_2D)175 gl.glBindTexture(gl.GL_TEXTURE_2D, self.gl_id)176 yield177 gl.glBindTexture(gl.GL_TEXTURE_2D, prev_id)178 179 def update(self, image):180 if image is not None:181 image = prepare_texture_data(image)182 assert self.is_compatible(image=image)183 with self.bind():184 fmt = get_texture_format(self.dtype, self.channels)185 gl.glPushClientAttrib(gl.GL_CLIENT_PIXEL_STORE_BIT)186 gl.glPixelStorei(gl.GL_UNPACK_ALIGNMENT, 1)187 gl.glTexImage2D(gl.GL_TEXTURE_2D, 0, fmt.internalformat, self.width, self.height, 0, fmt.format, fmt.type, image)188 if self.mipmap:189 gl.glGenerateMipmap(gl.GL_TEXTURE_2D)190 gl.glPopClientAttrib()191 192 def draw(self, *, pos=0, zoom=1, align=0, rint=False, color=1, alpha=1, rounding=0):193 zoom = np.broadcast_to(np.asarray(zoom, dtype='float32'), [2])194 size = zoom * [self.width, self.height]195 with self.bind():196 gl.glPushAttrib(gl.GL_ENABLE_BIT)197 gl.glEnable(gl.GL_TEXTURE_2D)198 draw_rect(pos=pos, size=size, align=align, rint=rint, color=color, alpha=alpha, rounding=rounding)199 gl.glPopAttrib()200 201 def is_compatible(self, *, image=None, width=None, height=None, channels=None, dtype=None): # pylint: disable=too-many-return-statements202 if image is not None:203 if image.ndim != 3:204 return False205 ih, iw, ic = image.shape206 if not self.is_compatible(width=iw, height=ih, channels=ic, dtype=image.dtype):207 return False208 if width is not None and self.width != width:209 return False210 if height is not None and self.height != height:211 return False212 if channels is not None and self.channels != channels:213 return False214 if dtype is not None and self.dtype != dtype:215 return False216 return True217 218#----------------------------------------------------------------------------219 220class Framebuffer:221 def __init__(self, *, texture=None, width=None, height=None, channels=None, dtype=None, msaa=0):222 self.texture = texture223 self.gl_id = None224 self.gl_color = None225 self.gl_depth_stencil = None226 self.msaa = msaa227 228 # Determine size and dtype.229 if texture is not None:230 assert isinstance(self.texture, Texture)231 self.width = texture.width232 self.height = texture.height233 self.channels = texture.channels234 self.dtype = texture.dtype235 else:236 assert width is not None and height is not None237 self.width = width238 self.height = height239 self.channels = channels if channels is not None else 4240 self.dtype = np.dtype(dtype) if dtype is not None else np.float32241 242 # Validate size and dtype.243 assert isinstance(self.width, int) and self.width >= 0244 assert isinstance(self.height, int) and self.height >= 0245 assert isinstance(self.channels, int) and self.channels >= 1246 assert width is None or width == self.width247 assert height is None or height == self.height248 assert channels is None or channels == self.channels249 assert dtype is None or dtype == self.dtype250 251 # Create framebuffer object.252 self.gl_id = gl.glGenFramebuffers(1)253 with self.bind():254 255 # Setup color buffer.256 if self.texture is not None:257 assert self.msaa == 0258 gl.glFramebufferTexture2D(gl.GL_FRAMEBUFFER, gl.GL_COLOR_ATTACHMENT0, gl.GL_TEXTURE_2D, self.texture.gl_id, 0)259 else:260 fmt = get_texture_format(self.dtype, self.channels)261 self.gl_color = gl.glGenRenderbuffers(1)262 gl.glBindRenderbuffer(gl.GL_RENDERBUFFER, self.gl_color)263 gl.glRenderbufferStorageMultisample(gl.GL_RENDERBUFFER, self.msaa, fmt.internalformat, self.width, self.height)264 gl.glFramebufferRenderbuffer(gl.GL_FRAMEBUFFER, gl.GL_COLOR_ATTACHMENT0, gl.GL_RENDERBUFFER, self.gl_color)265 266 # Setup depth/stencil buffer.267 self.gl_depth_stencil = gl.glGenRenderbuffers(1)268 gl.glBindRenderbuffer(gl.GL_RENDERBUFFER, self.gl_depth_stencil)269 gl.glRenderbufferStorageMultisample(gl.GL_RENDERBUFFER, self.msaa, gl.GL_DEPTH24_STENCIL8, self.width, self.height)270 gl.glFramebufferRenderbuffer(gl.GL_FRAMEBUFFER, gl.GL_DEPTH_STENCIL_ATTACHMENT, gl.GL_RENDERBUFFER, self.gl_depth_stencil)271 272 def delete(self):273 if self.gl_id is not None:274 gl.glDeleteFramebuffers([self.gl_id])275 self.gl_id = None276 if self.gl_color is not None:277 gl.glDeleteRenderbuffers(1, [self.gl_color])278 self.gl_color = None279 if self.gl_depth_stencil is not None:280 gl.glDeleteRenderbuffers(1, [self.gl_depth_stencil])281 self.gl_depth_stencil = None282 283 def __del__(self):284 try:285 self.delete()286 except:287 pass288 289 @contextlib.contextmanager290 def bind(self):291 prev_fbo = gl.glGetInteger(gl.GL_FRAMEBUFFER_BINDING)292 prev_rbo = gl.glGetInteger(gl.GL_RENDERBUFFER_BINDING)293 gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, self.gl_id)294 if self.width is not None and self.height is not None:295 gl.glViewport(0, 0, self.width, self.height)296 yield297 gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, prev_fbo)298 gl.glBindRenderbuffer(gl.GL_RENDERBUFFER, prev_rbo)299 300 def blit(self, dst=None):301 assert dst is None or isinstance(dst, Framebuffer)302 with self.bind():303 gl.glBindFramebuffer(gl.GL_DRAW_FRAMEBUFFER, 0 if dst is None else dst.fbo)304 gl.glBlitFramebuffer(0, 0, self.width, self.height, 0, 0, self.width, self.height, gl.GL_COLOR_BUFFER_BIT, gl.GL_NEAREST)305 306#----------------------------------------------------------------------------307 308def draw_shape(vertices, *, mode=gl.GL_TRIANGLE_FAN, pos=0, size=1, color=1, alpha=1):309 assert vertices.ndim == 2 and vertices.shape[1] == 2310 pos = np.broadcast_to(np.asarray(pos, dtype='float32'), [2])311 size = np.broadcast_to(np.asarray(size, dtype='float32'), [2])312 color = np.broadcast_to(np.asarray(color, dtype='float32'), [3])313 alpha = np.clip(np.broadcast_to(np.asarray(alpha, dtype='float32'), []), 0, 1)314 315 gl.glPushClientAttrib(gl.GL_CLIENT_VERTEX_ARRAY_BIT)316 gl.glPushAttrib(gl.GL_CURRENT_BIT | gl.GL_TRANSFORM_BIT)317 gl.glMatrixMode(gl.GL_MODELVIEW)318 gl.glPushMatrix()319 320 gl.glEnableClientState(gl.GL_VERTEX_ARRAY)321 gl.glEnableClientState(gl.GL_TEXTURE_COORD_ARRAY)322 gl.glVertexPointer(2, gl.GL_FLOAT, 0, vertices)323 gl.glTexCoordPointer(2, gl.GL_FLOAT, 0, vertices)324 gl.glTranslate(pos[0], pos[1], 0)325 gl.glScale(size[0], size[1], 1)326 gl.glColor4f(color[0] * alpha, color[1] * alpha, color[2] * alpha, alpha)327 gl.glDrawArrays(mode, 0, vertices.shape[0])328 329 gl.glPopMatrix()330 gl.glPopAttrib()331 gl.glPopClientAttrib()332 333#----------------------------------------------------------------------------334 335def draw_rect(*, pos=0, pos2=None, size=None, align=0, rint=False, color=1, alpha=1, rounding=0):336 assert pos2 is None or size is None337 pos = np.broadcast_to(np.asarray(pos, dtype='float32'), [2])338 pos2 = np.broadcast_to(np.asarray(pos2, dtype='float32'), [2]) if pos2 is not None else None339 size = np.broadcast_to(np.asarray(size, dtype='float32'), [2]) if size is not None else None340 size = size if size is not None else pos2 - pos if pos2 is not None else np.array([1, 1], dtype='float32')341 pos = pos - size * align342 if rint:343 pos = np.rint(pos)344 rounding = np.broadcast_to(np.asarray(rounding, dtype='float32'), [2])345 rounding = np.minimum(np.abs(rounding) / np.maximum(np.abs(size), 1e-8), 0.5)346 if np.min(rounding) == 0:347 rounding *= 0348 vertices = _setup_rect(float(rounding[0]), float(rounding[1]))349 draw_shape(vertices, mode=gl.GL_TRIANGLE_FAN, pos=pos, size=size, color=color, alpha=alpha)350 351@functools.lru_cache(maxsize=10000)352def _setup_rect(rx, ry):353 t = np.linspace(0, np.pi / 2, 1 if max(rx, ry) == 0 else 64)354 s = 1 - np.sin(t); c = 1 - np.cos(t)355 x = [c * rx, 1 - s * rx, 1 - c * rx, s * rx]356 y = [s * ry, c * ry, 1 - s * ry, 1 - c * ry]357 v = np.stack([x, y], axis=-1).reshape(-1, 2)358 return v.astype('float32')359 360#----------------------------------------------------------------------------361 362def draw_circle(*, center=0, radius=100, hole=0, color=1, alpha=1):363 hole = np.broadcast_to(np.asarray(hole, dtype='float32'), [])364 vertices = _setup_circle(float(hole))365 draw_shape(vertices, mode=gl.GL_TRIANGLE_STRIP, pos=center, size=radius, color=color, alpha=alpha)366 367@functools.lru_cache(maxsize=10000)368def _setup_circle(hole):369 t = np.linspace(0, np.pi * 2, 128)370 s = np.sin(t); c = np.cos(t)371 v = np.stack([c, s, c * hole, s * hole], axis=-1).reshape(-1, 2)372 return v.astype('float32')373 374#----------------------------------------------------------------------------375