hy3dgen/texgen/custom_rasterizer/custom_rasterizer/io_glb.py

# 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)