OpenMotionLab/MotionGPT
118
1"""Primitives, conforming to the glTF 2.0 standards as specified in2https://github.com/KhronosGroup/glTF/tree/master/specification/2.0#reference-primitive3 4Author: Matthew Matl5"""6import numpy as np7 8from OpenGL.GL import *9 10from .material import Material, MetallicRoughnessMaterial11from .constants import FLOAT_SZ, UINT_SZ, BufFlags, GLTF12from .utils import format_color_array13 14 15class Primitive(object):16 """A primitive object which can be rendered.17 18 Parameters19 ----------20 positions : (n, 3) float21 XYZ vertex positions.22 normals : (n, 3) float23 Normalized XYZ vertex normals.24 tangents : (n, 4) float25 XYZW vertex tangents where the w component is a sign value26 (either +1 or -1) indicating the handedness of the tangent basis.27 texcoord_0 : (n, 2) float28 The first set of UV texture coordinates.29 texcoord_1 : (n, 2) float30 The second set of UV texture coordinates.31 color_0 : (n, 4) float32 RGBA vertex colors.33 joints_0 : (n, 4) float34 Joint information.35 weights_0 : (n, 4) float36 Weight information for morphing.37 indices : (m, 3) int38 Face indices for triangle meshes or fans.39 material : :class:`Material`40 The material to apply to this primitive when rendering.41 mode : int42 The type of primitives to render, one of the following:43 44 - ``0``: POINTS45 - ``1``: LINES46 - ``2``: LINE_LOOP47 - ``3``: LINE_STRIP48 - ``4``: TRIANGLES49 - ``5``: TRIANGLES_STRIP50 - ``6``: TRIANGLES_FAN51 targets : (k,) int52 Morph target indices.53 poses : (x,4,4), float54 Array of 4x4 transformation matrices for instancing this object.55 """56 57 def __init__(self,58 positions,59 normals=None,60 tangents=None,61 texcoord_0=None,62 texcoord_1=None,63 color_0=None,64 joints_0=None,65 weights_0=None,66 indices=None,67 material=None,68 mode=None,69 targets=None,70 poses=None):71 72 if mode is None:73 mode = GLTF.TRIANGLES74 75 self.positions = positions76 self.normals = normals77 self.tangents = tangents78 self.texcoord_0 = texcoord_079 self.texcoord_1 = texcoord_180 self.color_0 = color_081 self.joints_0 = joints_082 self.weights_0 = weights_083 self.indices = indices84 self.material = material85 self.mode = mode86 self.targets = targets87 self.poses = poses88 89 self._bounds = None90 self._vaid = None91 self._buffers = []92 self._is_transparent = None93 self._buf_flags = None94 95 @property96 def positions(self):97 """(n,3) float : XYZ vertex positions.98 """99 return self._positions100 101 @positions.setter102 def positions(self, value):103 value = np.asanyarray(value, dtype=np.float32)104 self._positions = np.ascontiguousarray(value)105 self._bounds = None106 107 @property108 def normals(self):109 """(n,3) float : Normalized XYZ vertex normals.110 """111 return self._normals112 113 @normals.setter114 def normals(self, value):115 if value is not None:116 value = np.asanyarray(value, dtype=np.float32)117 value = np.ascontiguousarray(value)118 if value.shape != self.positions.shape:119 raise ValueError('Incorrect normals shape')120 self._normals = value121 122 @property123 def tangents(self):124 """(n,4) float : XYZW vertex tangents.125 """126 return self._tangents127 128 @tangents.setter129 def tangents(self, value):130 if value is not None:131 value = np.asanyarray(value, dtype=np.float32)132 value = np.ascontiguousarray(value)133 if value.shape != (self.positions.shape[0], 4):134 raise ValueError('Incorrect tangent shape')135 self._tangents = value136 137 @property138 def texcoord_0(self):139 """(n,2) float : The first set of UV texture coordinates.140 """141 return self._texcoord_0142 143 @texcoord_0.setter144 def texcoord_0(self, value):145 if value is not None:146 value = np.asanyarray(value, dtype=np.float32)147 value = np.ascontiguousarray(value)148 if (value.ndim != 2 or value.shape[0] != self.positions.shape[0] or149 value.shape[1] < 2):150 raise ValueError('Incorrect texture coordinate shape')151 if value.shape[1] > 2:152 value = value[:,:2]153 self._texcoord_0 = value154 155 @property156 def texcoord_1(self):157 """(n,2) float : The second set of UV texture coordinates.158 """159 return self._texcoord_1160 161 @texcoord_1.setter162 def texcoord_1(self, value):163 if value is not None:164 value = np.asanyarray(value, dtype=np.float32)165 value = np.ascontiguousarray(value)166 if (value.ndim != 2 or value.shape[0] != self.positions.shape[0] or167 value.shape[1] != 2):168 raise ValueError('Incorrect texture coordinate shape')169 self._texcoord_1 = value170 171 @property172 def color_0(self):173 """(n,4) float : RGBA vertex colors.174 """175 return self._color_0176 177 @color_0.setter178 def color_0(self, value):179 if value is not None:180 value = np.ascontiguousarray(181 format_color_array(value, shape=(len(self.positions), 4))182 )183 self._is_transparent = None184 self._color_0 = value185 186 @property187 def joints_0(self):188 """(n,4) float : Joint information.189 """190 return self._joints_0191 192 @joints_0.setter193 def joints_0(self, value):194 self._joints_0 = value195 196 @property197 def weights_0(self):198 """(n,4) float : Weight information for morphing.199 """200 return self._weights_0201 202 @weights_0.setter203 def weights_0(self, value):204 self._weights_0 = value205 206 @property207 def indices(self):208 """(m,3) int : Face indices for triangle meshes or fans.209 """210 return self._indices211 212 @indices.setter213 def indices(self, value):214 if value is not None:215 value = np.asanyarray(value, dtype=np.float32)216 value = np.ascontiguousarray(value)217 self._indices = value218 219 @property220 def material(self):221 """:class:`Material` : The material for this primitive.222 """223 return self._material224 225 @material.setter226 def material(self, value):227 # Create default material228 if value is None:229 value = MetallicRoughnessMaterial()230 else:231 if not isinstance(value, Material):232 raise TypeError('Object material must be of type Material')233 self._material = value234 235 @property236 def mode(self):237 """int : The type of primitive to render.238 """239 return self._mode240 241 @mode.setter242 def mode(self, value):243 value = int(value)244 if value < GLTF.POINTS or value > GLTF.TRIANGLE_FAN:245 raise ValueError('Invalid mode')246 self._mode = value247 248 @property249 def targets(self):250 """(k,) int : Morph target indices.251 """252 return self._targets253 254 @targets.setter255 def targets(self, value):256 self._targets = value257 258 @property259 def poses(self):260 """(x,4,4) float : Homogenous transforms for instancing this primitive.261 """262 return self._poses263 264 @poses.setter265 def poses(self, value):266 if value is not None:267 value = np.asanyarray(value, dtype=np.float32)268 value = np.ascontiguousarray(value)269 if value.ndim == 2:270 value = value[np.newaxis,:,:]271 if value.shape[1] != 4 or value.shape[2] != 4:272 raise ValueError('Pose matrices must be of shape (n,4,4), '273 'got {}'.format(value.shape))274 self._poses = value275 self._bounds = None276 277 @property278 def bounds(self):279 if self._bounds is None:280 self._bounds = self._compute_bounds()281 return self._bounds282 283 @property284 def centroid(self):285 """(3,) float : The centroid of the primitive's AABB.286 """287 return np.mean(self.bounds, axis=0)288 289 @property290 def extents(self):291 """(3,) float : The lengths of the axes of the primitive's AABB.292 """293 return np.diff(self.bounds, axis=0).reshape(-1)294 295 @property296 def scale(self):297 """(3,) float : The length of the diagonal of the primitive's AABB.298 """299 return np.linalg.norm(self.extents)300 301 @property302 def buf_flags(self):303 """int : The flags for the render buffer.304 """305 if self._buf_flags is None:306 self._buf_flags = self._compute_buf_flags()307 return self._buf_flags308 309 def delete(self):310 self._unbind()311 self._remove_from_context()312 313 @property314 def is_transparent(self):315 """bool : If True, the mesh is partially-transparent.316 """317 return self._compute_transparency()318 319 def _add_to_context(self):320 if self._vaid is not None:321 raise ValueError('Mesh is already bound to a context')322 323 # Generate and bind VAO324 self._vaid = glGenVertexArrays(1)325 glBindVertexArray(self._vaid)326 327 #######################################################################328 # Fill vertex buffer329 #######################################################################330 331 # Generate and bind vertex buffer332 vertexbuffer = glGenBuffers(1)333 self._buffers.append(vertexbuffer)334 glBindBuffer(GL_ARRAY_BUFFER, vertexbuffer)335 336 # positions337 vertex_data = self.positions338 attr_sizes = [3]339 340 # Normals341 if self.normals is not None:342 vertex_data = np.hstack((vertex_data, self.normals))343 attr_sizes.append(3)344 345 # Tangents346 if self.tangents is not None:347 vertex_data = np.hstack((vertex_data, self.tangents))348 attr_sizes.append(4)349 350 # Texture Coordinates351 if self.texcoord_0 is not None:352 vertex_data = np.hstack((vertex_data, self.texcoord_0))353 attr_sizes.append(2)354 if self.texcoord_1 is not None:355 vertex_data = np.hstack((vertex_data, self.texcoord_1))356 attr_sizes.append(2)357 358 # Color359 if self.color_0 is not None:360 vertex_data = np.hstack((vertex_data, self.color_0))361 attr_sizes.append(4)362 363 # TODO JOINTS AND WEIGHTS364 # PASS365 366 # Copy data to buffer367 vertex_data = np.ascontiguousarray(368 vertex_data.flatten().astype(np.float32)369 )370 glBufferData(371 GL_ARRAY_BUFFER, FLOAT_SZ * len(vertex_data),372 vertex_data, GL_STATIC_DRAW373 )374 total_sz = sum(attr_sizes)375 offset = 0376 for i, sz in enumerate(attr_sizes):377 glVertexAttribPointer(378 i, sz, GL_FLOAT, GL_FALSE, FLOAT_SZ * total_sz,379 ctypes.c_void_p(FLOAT_SZ * offset)380 )381 glEnableVertexAttribArray(i)382 offset += sz383 384 #######################################################################385 # Fill model matrix buffer386 #######################################################################387 388 if self.poses is not None:389 pose_data = np.ascontiguousarray(390 np.transpose(self.poses, [0,2,1]).flatten().astype(np.float32)391 )392 else:393 pose_data = np.ascontiguousarray(394 np.eye(4).flatten().astype(np.float32)395 )396 397 modelbuffer = glGenBuffers(1)398 self._buffers.append(modelbuffer)399 glBindBuffer(GL_ARRAY_BUFFER, modelbuffer)400 glBufferData(401 GL_ARRAY_BUFFER, FLOAT_SZ * len(pose_data),402 pose_data, GL_STATIC_DRAW403 )404 405 for i in range(0, 4):406 idx = i + len(attr_sizes)407 glEnableVertexAttribArray(idx)408 glVertexAttribPointer(409 idx, 4, GL_FLOAT, GL_FALSE, FLOAT_SZ * 4 * 4,410 ctypes.c_void_p(4 * FLOAT_SZ * i)411 )412 glVertexAttribDivisor(idx, 1)413 414 #######################################################################415 # Fill element buffer416 #######################################################################417 if self.indices is not None:418 elementbuffer = glGenBuffers(1)419 self._buffers.append(elementbuffer)420 glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, elementbuffer)421 glBufferData(GL_ELEMENT_ARRAY_BUFFER, UINT_SZ * self.indices.size,422 self.indices.flatten().astype(np.uint32),423 GL_STATIC_DRAW)424 425 glBindVertexArray(0)426 427 def _remove_from_context(self):428 if self._vaid is not None:429 glDeleteVertexArrays(1, [self._vaid])430 glDeleteBuffers(len(self._buffers), self._buffers)431 self._vaid = None432 self._buffers = []433 434 def _in_context(self):435 return self._vaid is not None436 437 def _bind(self):438 if self._vaid is None:439 raise ValueError('Cannot bind a Mesh that has not been added '440 'to a context')441 glBindVertexArray(self._vaid)442 443 def _unbind(self):444 glBindVertexArray(0)445 446 def _compute_bounds(self):447 """Compute the bounds of this object.448 """449 # Compute bounds of this object450 bounds = np.array([np.min(self.positions, axis=0),451 np.max(self.positions, axis=0)])452 453 # If instanced, compute translations for approximate bounds454 if self.poses is not None:455 bounds += np.array([np.min(self.poses[:,:3,3], axis=0),456 np.max(self.poses[:,:3,3], axis=0)])457 return bounds458 459 def _compute_transparency(self):460 """Compute whether or not this object is transparent.461 """462 if self.material.is_transparent:463 return True464 if self._is_transparent is None:465 self._is_transparent = False466 if self.color_0 is not None:467 if np.any(self._color_0[:,3] != 1.0):468 self._is_transparent = True469 return self._is_transparent470 471 def _compute_buf_flags(self):472 buf_flags = BufFlags.POSITION473 474 if self.normals is not None:475 buf_flags |= BufFlags.NORMAL476 if self.tangents is not None:477 buf_flags |= BufFlags.TANGENT478 if self.texcoord_0 is not None:479 buf_flags |= BufFlags.TEXCOORD_0480 if self.texcoord_1 is not None:481 buf_flags |= BufFlags.TEXCOORD_1482 if self.color_0 is not None:483 buf_flags |= BufFlags.COLOR_0484 if self.joints_0 is not None:485 buf_flags |= BufFlags.JOINTS_0486 if self.weights_0 is not None:487 buf_flags |= BufFlags.WEIGHTS_0488 489 return buf_flags490