Upload hy3dgen/texgen/custom_rasterizer/custom_rasterizer/io_glb.py with huggingface_hub

hy3dgen/texgen/custom_rasterizer/custom_rasterizer/io_glb.py

@@ -0,0 +1,248 @@
+# Open Source Model Licensed under the Apache License Version 2.0
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+
+# Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved.
+# The below software and/or models in this distribution may have been
+# modified by THL A29 Limited ("Tencent Modifications").
+# All Tencent Modifications are Copyright (C) THL A29 Limited.
+
+# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT
+# except for the third-party components listed below.
+# Hunyuan 3D does not impose any additional limitations beyond what is outlined
+# in the repsective licenses of these third-party components.
+# Users must comply with all terms and conditions of original licenses of these third-party
+# components and must ensure that the usage of the third party components adheres to
+# all relevant laws and regulations.
+
+# For avoidance of doubts, Hunyuan 3D means the large language models and
+# their software and algorithms, including trained model weights, parameters (including
+# optimizer states), machine-learning model code, inference-enabling code, training-enabling code,
+# fine-tuning enabling code and other elements of the foregoing made publicly available
+# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
+
+import base64
+import io
+import os
+
+import numpy as np
+from PIL import Image as PILImage
+from pygltflib import GLTF2
+from scipy.spatial.transform import Rotation as R
+
+
+# Function to extract buffer data
+def get_buffer_data(gltf, buffer_view):
+    buffer = gltf.buffers[buffer_view.buffer]
+    buffer_data = gltf.get_data_from_buffer_uri(buffer.uri)
+    byte_offset = buffer_view.byteOffset if buffer_view.byteOffset else 0
+    byte_length = buffer_view.byteLength
+    return buffer_data[byte_offset:byte_offset + byte_length]
+
+
+# Function to extract attribute data
+def get_attribute_data(gltf, accessor_index):
+    accessor = gltf.accessors[accessor_index]
+    buffer_view = gltf.bufferViews[accessor.bufferView]
+    buffer_data = get_buffer_data(gltf, buffer_view)
+
+    comptype = {5120: np.int8, 5121: np.uint8, 5122: np.int16, 5123: np.uint16, 5125: np.uint32, 5126: np.float32}
+    dtype = comptype[accessor.componentType]
+
+    t2n = {'SCALAR': 1, 'VEC2': 2, 'VEC3': 3, 'VEC4': 4, 'MAT2': 4, 'MAT3': 9, 'MAT4': 16}
+    num_components = t2n[accessor.type]
+
+    # Calculate the correct slice of data
+    byte_offset = accessor.byteOffset if accessor.byteOffset else 0
+    byte_stride = buffer_view.byteStride if buffer_view.byteStride else num_components * np.dtype(dtype).itemsize
+    count = accessor.count
+
+    # Extract the attribute data
+    attribute_data = np.zeros((count, num_components), dtype=dtype)
+    for i in range(count):
+        start = byte_offset + i * byte_stride
+        end = start + num_components * np.dtype(dtype).itemsize
+        attribute_data[i] = np.frombuffer(buffer_data[start:end], dtype=dtype)
+
+    return attribute_data
+
+
+# Function to extract image data
+def get_image_data(gltf, image, folder):
+    if image.uri:
+        if image.uri.startswith('data:'):
+            # Data URI
+            header, encoded = image.uri.split(',', 1)
+            data = base64.b64decode(encoded)
+        else:
+            # External file
+            fn = image.uri
+            if not os.path.isabs(fn):
+                fn = folder + '/' + fn
+            with open(fn, 'rb') as f:
+                data = f.read()
+    else:
+        buffer_view = gltf.bufferViews[image.bufferView]
+        data = get_buffer_data(gltf, buffer_view)
+    return data
+
+
+# Function to convert triangle strip to triangles
+def convert_triangle_strip_to_triangles(indices):
+    triangles = []
+    for i in range(len(indices) - 2):
+        if i % 2 == 0:
+            triangles.append([indices[i], indices[i + 1], indices[i + 2]])
+        else:
+            triangles.append([indices[i], indices[i + 2], indices[i + 1]])
+    return np.array(triangles).reshape(-1, 3)
+
+
+# Function to convert triangle fan to triangles
+def convert_triangle_fan_to_triangles(indices):
+    triangles = []
+    for i in range(1, len(indices) - 1):
+        triangles.append([indices[0], indices[i], indices[i + 1]])
+    return np.array(triangles).reshape(-1, 3)
+
+
+# Function to get the transformation matrix from a node
+def get_node_transform(node):
+    if node.matrix:
+        return np.array(node.matrix).reshape(4, 4).T
+    else:
+        T = np.eye(4)
+        if node.translation:
+            T[:3, 3] = node.translation
+        if node.rotation:
+            R_mat = R.from_quat(node.rotation).as_matrix()
+            T[:3, :3] = R_mat
+        if node.scale:
+            S = np.diag(node.scale + [1])
+            T = T @ S
+        return T
+
+
+def get_world_transform(gltf, node_index, parents, world_transforms):
+    if parents[node_index] == -2:
+        return world_transforms[node_index]
+
+    node = gltf.nodes[node_index]
+    if parents[node_index] == -1:
+        world_transforms[node_index] = get_node_transform(node)
+        parents[node_index] = -2
+        return world_transforms[node_index]
+
+    parent_index = parents[node_index]
+    parent_transform = get_world_transform(gltf, parent_index, parents, world_transforms)
+    world_transforms[node_index] = parent_transform @ get_node_transform(node)
+    parents[node_index] = -2
+    return world_transforms[node_index]
+
+
+def LoadGlb(path):
+    # Load the GLB file using pygltflib
+    gltf = GLTF2().load(path)
+
+    primitives = []
+    images = {}
+    # Iterate through the meshes in the GLB file
+
+    world_transforms = [np.identity(4) for i in range(len(gltf.nodes))]
+    parents = [-1 for i in range(len(gltf.nodes))]
+    for node_index, node in enumerate(gltf.nodes):
+        for idx in node.children:
+            parents[idx] = node_index
+    # for i in range(len(gltf.nodes)):
+    #    get_world_transform(gltf, i, parents, world_transform)
+
+    for node_index, node in enumerate(gltf.nodes):
+        if node.mesh is not None:
+            world_transform = get_world_transform(gltf, node_index, parents, world_transforms)
+            # Iterate through the primitives in the mesh
+            mesh = gltf.meshes[node.mesh]
+            for primitive in mesh.primitives:
+                # Access the attributes of the primitive
+                attributes = primitive.attributes.__dict__
+                mode = primitive.mode if primitive.mode is not None else 4  # Default to TRIANGLES
+                result = {}
+                if primitive.indices is not None:
+                    indices = get_attribute_data(gltf, primitive.indices)
+                    if mode == 4:  # TRIANGLES
+                        face_indices = indices.reshape(-1, 3)
+                    elif mode == 5:  # TRIANGLE_STRIP
+                        face_indices = convert_triangle_strip_to_triangles(indices)
+                    elif mode == 6:  # TRIANGLE_FAN
+                        face_indices = convert_triangle_fan_to_triangles(indices)
+                    else:
+                        continue
+                    result['F'] = face_indices
+
+                # Extract vertex positions
+                if 'POSITION' in attributes and attributes['POSITION'] is not None:
+                    positions = get_attribute_data(gltf, attributes['POSITION'])
+                    # Apply the world transformation to the positions
+                    positions_homogeneous = np.hstack([positions, np.ones((positions.shape[0], 1))])
+                    transformed_positions = (world_transform @ positions_homogeneous.T).T[:, :3]
+                    result['V'] = transformed_positions
+
+                # Extract vertex colors
+                if 'COLOR_0' in attributes and attributes['COLOR_0'] is not None:
+                    colors = get_attribute_data(gltf, attributes['COLOR_0'])
+                    if colors.shape[-1] > 3:
+                        colors = colors[..., :3]
+                    result['VC'] = colors
+
+                # Extract UVs
+                if 'TEXCOORD_0' in attributes and not attributes['TEXCOORD_0'] is None:
+                    uvs = get_attribute_data(gltf, attributes['TEXCOORD_0'])
+                    result['UV'] = uvs
+
+                if primitive.material is not None:
+                    material = gltf.materials[primitive.material]
+                    if material.pbrMetallicRoughness is not None and material.pbrMetallicRoughness.baseColorTexture is not None:
+                        texture_index = material.pbrMetallicRoughness.baseColorTexture.index
+                        texture = gltf.textures[texture_index]
+                        image_index = texture.source
+                        if not image_index in images:
+                            image = gltf.images[image_index]
+                            image_data = get_image_data(gltf, image, os.path.dirname(path))
+                            pil_image = PILImage.open(io.BytesIO(image_data))
+                            if pil_image.mode != 'RGB':
+                                pil_image = pil_image.convert('RGB')
+                            images[image_index] = pil_image
+                        result['TEX'] = image_index
+                    elif material.emissiveTexture is not None:
+                        texture_index = material.emissiveTexture.index
+                        texture = gltf.textures[texture_index]
+                        image_index = texture.source
+                        if not image_index in images:
+                            image = gltf.images[image_index]
+                            image_data = get_image_data(gltf, image, os.path.dirname(path))
+                            pil_image = PILImage.open(io.BytesIO(image_data))
+                            if pil_image.mode != 'RGB':
+                                pil_image = pil_image.convert('RGB')
+                            images[image_index] = pil_image
+                        result['TEX'] = image_index
+                    else:
+                        if material.pbrMetallicRoughness is not None:
+                            base_color = material.pbrMetallicRoughness.baseColorFactor
+                        else:
+                            base_color = np.array([0.8, 0.8, 0.8], dtype=np.float32)
+                        result['MC'] = base_color
+
+                primitives.append(result)
+
+    return primitives, images
+
+
+def RotatePrimitives(primitives, transform):
+    for i in range(len(primitives)):
+        if 'V' in primitives[i]:
+            primitives[i]['V'] = primitives[i]['V'] @ transform.T
+
+
+if __name__ == '__main__':
+    path = 'data/test.glb'
+    LoadGlb(path)