UAV/tools/merge_two_obj_simple.py

245 lines
9.1 KiB
Python
Raw Permalink Normal View History

2025-01-02 10:36:16 +08:00
import os
import shutil
def read_obj(file_path):
"""读取OBJ文件返回顶点、纹理坐标、法线、面的列表和MTL文件名"""
vertices = [] # v
tex_coords = [] # vt
normals = [] # vn
faces = [] # f
face_materials = [] # 每个面对应的材质名称
mtl_file = None # mtl文件名
current_material = None # 当前使用的材质
with open(file_path, 'r') as file:
for line in file:
if line.startswith('#') or not line.strip():
continue
parts = line.strip().split()
if not parts:
continue
if parts[0] == 'mtllib': # MTL文件引用
mtl_file = parts[1]
elif parts[0] == 'usemtl': # 材质使用
current_material = parts[1]
elif parts[0] == 'v': # 顶点
vertices.append([float(parts[1]), float(parts[2]), float(parts[3])])
elif parts[0] == 'vt': # 纹理坐标
tex_coords.append([float(parts[1]), float(parts[2])])
elif parts[0] == 'vn': # 法线
normals.append([float(parts[1]), float(parts[2]), float(parts[3])])
elif parts[0] == 'f': # 面
# 处理面的索引 (v/vt/vn)
face_v = []
face_vt = []
face_vn = []
for p in parts[1:]:
indices = p.split('/')
face_v.append(int(indices[0]))
if len(indices) > 1 and indices[1]:
face_vt.append(int(indices[1]))
if len(indices) > 2:
face_vn.append(int(indices[2]))
faces.append((face_v, face_vt, face_vn))
face_materials.append(current_material) # 记录这个面使用的材质
return vertices, tex_coords, normals, faces, face_materials, mtl_file
def read_mtl(mtl_path: str) -> tuple:
"""读取MTL文件内容
Returns:
tuple: (文件内容列表, 材质名称列表)
"""
content = []
material_names = []
with open(mtl_path, 'r') as f:
content = f.readlines()
for line in content:
if line.startswith('newmtl'):
material_names.append(line.strip().split()[1])
return content, material_names
def update_mtl_content(content: list, model_id: int) -> tuple:
"""更新MTL文件内容修改材质名称和纹理文件路径
Returns:
tuple: (更新后的内容, 更新后的材质名称列表)
"""
updated_lines = []
updated_material_names = []
current_material = None
for line in content:
if line.startswith('newmtl'):
# 为材质名称添加前缀
parts = line.strip().split()
material_name = parts[1]
current_material = f"grid_{model_id}_{material_name}"
updated_material_names.append(current_material)
updated_lines.append(f"newmtl {current_material}\n")
elif line.startswith('map_'):
# 更新纹理文件路径
parts = line.strip().split()
texture_file = os.path.basename(parts[-1])
parts[-1] = f"grid_{model_id}_{texture_file}"
updated_lines.append(' '.join(parts) + '\n')
else:
updated_lines.append(line)
return updated_lines, updated_material_names
def merge_mtl_files(mtl1_path: str, mtl2_path: str, output_path: str) -> tuple:
"""合并两个MTL文件
Returns:
tuple: (第一个模型的材质名称列表, 第二个模型的材质名称列表)
"""
# 读取两个MTL文件
content1, materials1 = read_mtl(mtl1_path)
content2, materials2 = read_mtl(mtl2_path)
# 更新两个MTL的内容
updated_content1, updated_materials1 = update_mtl_content(content1, 0)
updated_content2, updated_materials2 = update_mtl_content(content2, 1)
# 合并并写入新的MTL文件
with open(output_path, 'w') as f:
f.writelines(updated_content1)
f.write('\n') # 添加分隔行
f.writelines(updated_content2)
return updated_materials1, updated_materials2
def write_obj(file_path, vertices, tex_coords, normals, faces, face_materials, mtl_file=None):
"""将顶点、纹理坐标、法线和面写入到OBJ文件"""
with open(file_path, 'w') as file:
# 写入MTL文件引用
if mtl_file:
file.write(f"mtllib {mtl_file}\n")
# 写入顶点
for v in vertices:
file.write(f"v {v[0]} {v[1]} {v[2]}\n")
# 写入纹理坐标
for vt in tex_coords:
file.write(f"vt {vt[0]} {vt[1]}\n")
# 写入法线
for vn in normals:
file.write(f"vn {vn[0]} {vn[1]} {vn[2]}\n")
# 写入面(按材质分组)
current_material = None
for face, material in zip(faces, face_materials):
# 如果材质发生变化写入新的usemtl
if material != current_material:
file.write(f"usemtl {material}\n")
current_material = material
face_str = "f"
for j in range(len(face[0])):
face_str += " "
face_str += str(face[0][j])
if face[1]:
face_str += f"/{face[1][j]}"
else:
face_str += "/"
if face[2]:
face_str += f"/{face[2][j]}"
else:
face_str += "/"
file.write(face_str + "\n")
def translate_vertices(vertices, translation):
"""平移顶点"""
return [[v[0] + translation[0], v[1] + translation[1], v[2] + translation[2]] for v in vertices]
def copy_mtl_and_textures(src_dir: str, dst_dir: str, model_id: int):
"""复制MTL文件和相关的纹理文件并重命名避免冲突
Args:
src_dir: 源目录包含MTL和纹理文件
dst_dir: 目标目录
model_id: 模型ID用于重命名
"""
# 复制并重命名纹理文件
for file in os.listdir(src_dir):
if file.lower().endswith('.png'):
src_file = os.path.join(src_dir, file)
new_name = f"grid_{model_id}_{file}"
dst_file = os.path.join(dst_dir, new_name)
shutil.copy2(src_file, dst_file)
print(f"复制纹理文件: {file} -> {new_name}")
def merge_objs(obj1_path, obj2_path, output_path):
"""合并两个OBJ文件"""
print(f"开始合并OBJ模型:\n输入1: {obj1_path}\n输入2: {obj2_path}")
# 读取两个obj文件
vertices1, tex_coords1, normals1, faces1, face_materials1, mtl1 = read_obj(obj1_path)
vertices2, tex_coords2, normals2, faces2, face_materials2, mtl2 = read_obj(obj2_path)
# 固定平移量(0, 1000, 0)
translation = (0, 1000, 0)
# 平移第二个模型的顶点
vertices2_translated = translate_vertices(vertices2, translation)
# 计算偏移量
v_offset = len(vertices1)
vt_offset = len(tex_coords1)
vn_offset = len(normals1)
# 合并顶点、纹理坐标和法线
all_vertices = vertices1 + vertices2_translated
all_tex_coords = tex_coords1 + tex_coords2
all_normals = normals1 + normals2
# 调整第二个模型的面索引和材质名称
all_faces = faces1.copy()
all_face_materials = face_materials1.copy()
for face, material in zip(faces2, face_materials2):
new_face_v = [f + v_offset for f in face[0]]
new_face_vt = [f + vt_offset for f in face[1]] if face[1] else []
new_face_vn = [f + vn_offset for f in face[2]] if face[2] else []
all_faces.append((new_face_v, new_face_vt, new_face_vn))
# 为第二个模型的材质名称添加前缀
all_face_materials.append(f"grid_1_{material}")
# 为第一个模型的材质添加前缀
all_face_materials[:len(faces1)] = [f"grid_0_{mat}" for mat in face_materials1]
# 创建输出子目录
output_dir = os.path.dirname(output_path)
os.makedirs(output_dir, exist_ok=True)
# 复制并重命名两个模型的纹理文件
src_dir1 = os.path.dirname(obj1_path)
src_dir2 = os.path.dirname(obj2_path)
copy_mtl_and_textures(src_dir1, output_dir, 0)
copy_mtl_and_textures(src_dir2, output_dir, 1)
# 合并MTL文件并获取材质名称
src_mtl1 = os.path.join(src_dir1, mtl1)
src_mtl2 = os.path.join(src_dir2, mtl2)
dst_mtl = os.path.join(output_dir, "merged_model.mtl")
merge_mtl_files(src_mtl1, src_mtl2, dst_mtl)
# 写入合并后的obj文件
write_obj(output_path, all_vertices, all_tex_coords, all_normals,
all_faces, all_face_materials, "merged_model.mtl")
print(f"模型合并成功,已保存至: {output_path}")
if __name__ == "__main__":
# 测试参数
obj1_path = r"G:\ODM_output\1009\grid_0_0\project\odm_texturing\odm_textured_model_geo.obj"
obj2_path = r"G:\ODM_output\1009\grid_0_1\project\odm_texturing\odm_textured_model_geo.obj"
output_dir = r"G:\ODM_output\1009\merge_test"
# 创建输出目录
os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(output_dir, "merged_test.obj")
# 执行合并
merge_objs(obj1_path, obj2_path, output_path)