init

hy3dpaint/convert_utils.py

@@ -0,0 +1,140 @@
+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和roughness贴图合并为一张贴图
+    GLB格式要求metallic在B通道，roughness在G通道
+    """
+    # 加载贴图
+    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):
+    """
+    使用pygltflib创建包含完整PBR材质的GLB文件
+
+    textures_dict = {
+        'albedo': 'path/to/albedo.png',
+        'metallic': 'path/to/metallic.png',
+        'roughness': 'path/to/roughness.png',
+        'normal': 'path/to/normal.png',  # 可选
+        'ao': 'path/to/ao.png'  # 可选
+    }
+    """
+    # 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):
+        """将图像转换为data URI"""
+        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}")
+
+