MaterialMVP/convert_utils.py at main · ZebinHe/MaterialMVP · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
import trimesh
import pygltflib
import numpy as np
from PIL import Image
import base64
import io


def combine_metallic_roughness(metallic_path, roughness_path, output_path):

    # 加载贴图
    metallic_img = Image.open(metallic_path).convert("L")  # 转为灰度
    roughness_img = Image.open(roughness_path).convert("L")  # 转为灰度

    # 确保尺寸一致
    if metallic_img.size != roughness_img.size:
        roughness_img = roughness_img.resize(metallic_img.size)

    # 创建RGB图像
    width, height = metallic_img.size
    combined = Image.new("RGB", (width, height))

    # 转为numpy数组便于操作
    metallic_array = np.array(metallic_img)
    roughness_array = np.array(roughness_img)

    # 创建合并的数组 (R, G, B) = (AO, Roughness, Metallic)
    combined_array = np.zeros((height, width, 3), dtype=np.uint8)
    combined_array[:, :, 0] = 255  # R通道：AO (如果没有AO贴图，设为白色)
    combined_array[:, :, 1] = roughness_array  # G通道：Roughness
    combined_array[:, :, 2] = metallic_array  # B通道：Metallic

    # 转回PIL图像并保存
    combined = Image.fromarray(combined_array)
    combined.save(output_path)
    return output_path


def create_glb_with_pbr_materials(obj_path, textures_dict, output_path):

    # 1. 加载OBJ文件
    mesh = trimesh.load(obj_path)

    # 2. 先导出为临时GLB
    temp_glb = "temp.glb"
    mesh.export(temp_glb)

    # 3. 加载GLB文件进行材质编辑
    gltf = pygltflib.GLTF2().load(temp_glb)

    # 4. 准备纹理数据
    def image_to_data_uri(image_path):

        with open(image_path, "rb") as f:
            image_data = f.read()
        encoded = base64.b64encode(image_data).decode()
        return f"data:image/png;base64,{encoded}"

    # 5. 合并metallic和roughness
    if "metallic" in textures_dict and "roughness" in textures_dict:
        mr_combined_path = "mr_combined.png"
        combine_metallic_roughness(textures_dict["metallic"], textures_dict["roughness"], mr_combined_path)
        textures_dict["metallicRoughness"] = mr_combined_path

    # 6. 添加图像到GLTF
    images = []
    textures = []

    texture_mapping = {
        "albedo": "baseColorTexture",
        "metallicRoughness": "metallicRoughnessTexture",
        "normal": "normalTexture",
        "ao": "occlusionTexture",
    }

    for tex_type, tex_path in textures_dict.items():
        if tex_type in texture_mapping and tex_path:
            # 添加图像
            image = pygltflib.Image(uri=image_to_data_uri(tex_path))
            images.append(image)

            # 添加纹理
            texture = pygltflib.Texture(source=len(images) - 1)
            textures.append(texture)

    # 7. 创建PBR材质
    pbr_metallic_roughness = pygltflib.PbrMetallicRoughness(
        baseColorFactor=[1.0, 1.0, 1.0, 1.0], metallicFactor=1.0, roughnessFactor=1.0
    )

    # 设置纹理索引
    texture_index = 0
    if "albedo" in textures_dict:
        pbr_metallic_roughness.baseColorTexture = pygltflib.TextureInfo(index=texture_index)
        texture_index += 1

    if "metallicRoughness" in textures_dict:
        pbr_metallic_roughness.metallicRoughnessTexture = pygltflib.TextureInfo(index=texture_index)
        texture_index += 1

    # 创建材质
    material = pygltflib.Material(name="PBR_Material", pbrMetallicRoughness=pbr_metallic_roughness)

    # 添加法线贴图
    if "normal" in textures_dict:
        material.normalTexture = pygltflib.NormalTextureInfo(index=texture_index)
        texture_index += 1

    # 添加AO贴图
    if "ao" in textures_dict:
        material.occlusionTexture = pygltflib.OcclusionTextureInfo(index=texture_index)

    # 8. 更新GLTF
    gltf.images = images
    gltf.textures = textures
    gltf.materials = [material]

    # 确保mesh使用材质
    if gltf.meshes:
        for primitive in gltf.meshes[0].primitives:
            primitive.material = 0

    # 9. 保存最终GLB
    gltf.save(output_path)
    print(f"PBR GLB文件已保存: {output_path}")


if __name__ == "__main__":
    create_glb_with_pbr_materials(
        "/root/save_dir/d4c88be8-4761-434d-a1b5-92baed95546c/textured_mesh.obj",
        {
            "albedo": "/root/save_dir/d4c88be8-4761-434d-a1b5-92baed95546c/textured_mesh.jpg",
            "metallic": "/root/save_dir/d4c88be8-4761-434d-a1b5-92baed95546c/textured_mesh_metallic.jpg",
            "roughness": "/root/save_dir/d4c88be8-4761-434d-a1b5-92baed95546c/textured_mesh_roughness.jpg",
        },
        "/root/save_dir/d4c88be8-4761-434d-a1b5-92baed95546c/textured_mesh.glb",
    )