388 lines
15 KiB
Python
388 lines
15 KiB
Python
import os
|
||
import logging
|
||
import numpy as np
|
||
from typing import Dict
|
||
import pandas as pd
|
||
import shutil
|
||
import time
|
||
|
||
|
||
class MergeObj:
|
||
def __init__(self, output_dir: str):
|
||
self.output_dir = output_dir
|
||
self.logger = logging.getLogger('UAV_Preprocess.MergeObj')
|
||
|
||
def read_obj(self, 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 write_obj(self, 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 i in range(len(face[0])):
|
||
face_str += " "
|
||
face_str += str(face[0][i])
|
||
if face[1]:
|
||
face_str += f"/{face[1][i]}"
|
||
else:
|
||
face_str += "/"
|
||
if face[2]:
|
||
face_str += f"/{face[2][i]}"
|
||
else:
|
||
face_str += "/"
|
||
file.write(face_str + "\n")
|
||
|
||
def translate_vertices(self, vertices, translation):
|
||
"""平移顶点"""
|
||
return [[v[0] + translation[0], v[1] + translation[1], v[2] + translation[2]] for v in vertices]
|
||
|
||
def merge_two_objs(self, obj1_path: str, obj2_path: str, output_path: str, translation, grid_id1: tuple, grid_id2: tuple):
|
||
"""合并两个OBJ文件"""
|
||
try:
|
||
self.logger.info(f"开始合并OBJ模型:\n输入1: {obj1_path}\n输入2: {obj2_path}")
|
||
|
||
# 读取两个obj文件
|
||
vertices1, tex_coords1, normals1, faces1, face_materials1, mtl1 = self.read_obj(
|
||
obj1_path)
|
||
vertices2, tex_coords2, normals2, faces2, face_materials2, mtl2 = self.read_obj(
|
||
obj2_path)
|
||
|
||
# 读取MTL文件内容以获取正确的材质名称
|
||
src_dir1 = os.path.dirname(obj1_path)
|
||
src_dir2 = os.path.dirname(obj2_path)
|
||
mtl1_path = os.path.join(src_dir1, mtl1)
|
||
mtl2_path = os.path.join(src_dir2, mtl2)
|
||
|
||
# 读取并更新材质内容
|
||
materials1 = self.read_mtl(mtl1_path)
|
||
materials2 = self.read_mtl(mtl2_path)
|
||
|
||
# 创建材质名称映射（使用与MTL文件相同的命名格式）
|
||
material_map1 = {}
|
||
material_map2 = {}
|
||
|
||
# 处理第一个模型的材质映射
|
||
for old_name in materials1.keys():
|
||
if "grid_0_0" in obj1_path:
|
||
material_map1[old_name] = f"material_{grid_id1[0]}_{grid_id1[1]}_{old_name}"
|
||
else:
|
||
# 更新完一次后，之后就不用再更新了
|
||
material_map1[old_name] = old_name
|
||
|
||
# 处理第二个模型的材质映射
|
||
for old_name in materials2.keys():
|
||
material_map2[old_name] = f"material_{grid_id2[0]}_{grid_id2[1]}_{old_name}"
|
||
|
||
# 平移第二个模型的顶点
|
||
vertices2_translated = self.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 = []
|
||
|
||
# 更新第一个模型的材质名称
|
||
for material in face_materials1:
|
||
all_face_materials.append(material_map1.get(material))
|
||
|
||
# 更新第二个模型的面索引和材质名称
|
||
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(material_map2.get(material))
|
||
|
||
# 写入合并后的obj文件，使用与MTL文件相同的名称
|
||
mtl_filename = "merged_model.mtl" # 使用固定的MTL文件名
|
||
self.write_obj(output_path, all_vertices, all_tex_coords, all_normals,
|
||
all_faces, all_face_materials, mtl_filename)
|
||
self.logger.info(f"模型合并成功，已保存至: {output_path}")
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"合并OBJ模型时发生错误: {str(e)}", exc_info=True)
|
||
raise
|
||
|
||
def read_mtl(self, mtl_path: str) -> dict:
|
||
"""读取MTL文件内容
|
||
Returns:
|
||
dict: 材质名称到材质信息的映射
|
||
"""
|
||
materials = {}
|
||
current_material = None
|
||
|
||
with open(mtl_path, 'r') as f:
|
||
content = f.read()
|
||
for line in content.strip().split('\n'):
|
||
if not line:
|
||
continue
|
||
|
||
parts = line.split()
|
||
if not parts:
|
||
continue
|
||
|
||
if parts[0] == 'newmtl':
|
||
current_material = parts[1]
|
||
materials[current_material] = []
|
||
elif current_material:
|
||
materials[current_material].append(line)
|
||
|
||
return materials
|
||
|
||
def copy_and_rename_texture(self, src_dir: str, dst_dir: str, grid_id: tuple) -> dict:
|
||
"""复制并重命名纹理文件
|
||
Args:
|
||
src_dir: 源纹理目录
|
||
dst_dir: 目标纹理目录
|
||
grid_id: 网格ID
|
||
Returns:
|
||
dict: 原始文件名到新文件名的映射
|
||
"""
|
||
texture_map = {}
|
||
os.makedirs(dst_dir, exist_ok=True)
|
||
|
||
for file in os.listdir(src_dir):
|
||
if file.lower().endswith(('.png', '.jpg', '.jpeg')):
|
||
# 生成新的文件名：grid_0_1_texture.png
|
||
new_name = f"grid_{grid_id[0]}_{grid_id[1]}_{file}"
|
||
src_path = os.path.join(src_dir, file)
|
||
dst_path = os.path.join(dst_dir, new_name)
|
||
|
||
# 复制文件
|
||
shutil.copy2(src_path, dst_path)
|
||
texture_map[file] = new_name
|
||
self.logger.debug(f"复制纹理文件: {file} -> {new_name}")
|
||
|
||
return texture_map
|
||
|
||
def update_mtl_content(self, materials: dict, texture_map: dict, grid_id: tuple) -> dict:
|
||
"""更新材质内容，修改材质名称和纹理路径
|
||
Args:
|
||
materials: 原始材质信息
|
||
texture_map: 纹理文件映射
|
||
grid_id: 网格ID
|
||
Returns:
|
||
dict: 更新后的材质信息
|
||
"""
|
||
updated_materials = {}
|
||
|
||
for mat_name, content in materials.items():
|
||
# 为材质名称添加网格ID前缀，与OBJ文件中的usemtl保持一致
|
||
new_mat_name = f"material_{grid_id[0]}_{grid_id[1]}_{mat_name}"
|
||
updated_content = []
|
||
|
||
for line in content:
|
||
if line.startswith('map_'): # 更新纹理文件路径
|
||
parts = line.split()
|
||
old_texture = parts[-1]
|
||
if old_texture in texture_map:
|
||
parts[-1] = texture_map[old_texture]
|
||
line = ' '.join(parts)
|
||
updated_content.append(line)
|
||
|
||
updated_materials[new_mat_name] = updated_content
|
||
|
||
return updated_materials
|
||
|
||
def merge_grid_obj(self, grid_points: Dict[tuple, pd.DataFrame], translations: Dict[tuple, tuple]):
|
||
"""合并所有网格的OBJ模型"""
|
||
if len(grid_points) == 1:
|
||
self.logger.info("只有一个网格，无需合并")
|
||
return
|
||
|
||
try:
|
||
# 创建输出目录
|
||
output_model_dir = os.path.join(self.output_dir, "merged_model")
|
||
os.makedirs(output_model_dir, exist_ok=True)
|
||
|
||
# 获取所有有效的网格文件
|
||
grid_files = {}
|
||
for grid_id, points in grid_points.items():
|
||
base_dir = os.path.join(
|
||
self.output_dir,
|
||
f"grid_{grid_id[0]}_{grid_id[1]}",
|
||
"project",
|
||
"odm_texturing"
|
||
)
|
||
obj_path = os.path.join(base_dir, "odm_textured_model_geo.obj")
|
||
mtl_path = os.path.join(base_dir, "odm_textured_model_geo.mtl")
|
||
|
||
if not os.path.exists(obj_path) or not os.path.exists(mtl_path):
|
||
self.logger.warning(
|
||
f"网格 ({grid_id[0]},{grid_id[1]}) 的文件不存在")
|
||
continue
|
||
|
||
grid_files[grid_id] = {
|
||
'obj': obj_path,
|
||
'mtl': mtl_path,
|
||
'dir': base_dir
|
||
}
|
||
|
||
if not grid_files:
|
||
self.logger.error("没有找到有效的文件")
|
||
return
|
||
|
||
# 收集所有材质和纹理信息
|
||
all_materials = {}
|
||
for grid_id, files in grid_files.items():
|
||
# 复制并重命名纹理文件
|
||
texture_map = self.copy_and_rename_texture(
|
||
files['dir'],
|
||
output_model_dir,
|
||
grid_id
|
||
)
|
||
|
||
# 读取并更新MTL内容
|
||
materials = self.read_mtl(files['mtl'])
|
||
updated_materials = self.update_mtl_content(
|
||
materials,
|
||
texture_map,
|
||
grid_id
|
||
)
|
||
all_materials.update(updated_materials)
|
||
|
||
# 写入合并后的MTL文件
|
||
final_mtl = os.path.join(output_model_dir, "merged_model.mtl")
|
||
with open(final_mtl, 'w') as f:
|
||
for mat_name, content in all_materials.items():
|
||
f.write(f"newmtl {mat_name}\n")
|
||
for line in content:
|
||
f.write(f"{line}\n")
|
||
f.write("\n")
|
||
|
||
# 合并OBJ文件
|
||
reference_id = list(grid_files.keys())[0]
|
||
merged_obj = grid_files[reference_id]['obj']
|
||
temp_files = [] # 记录所有中间文件
|
||
|
||
for grid_id, files in list(grid_files.items())[1:]:
|
||
translation = translations[grid_id]
|
||
translation = (translation[0], translation[1], 0)
|
||
|
||
# 生成临时输出文件名
|
||
temp_output = os.path.join(
|
||
output_model_dir,
|
||
f"temp_merged_{int(time.time())}.obj"
|
||
)
|
||
temp_files.append(temp_output) # 添加到临时文件列表
|
||
|
||
self.merge_two_objs(
|
||
merged_obj, files['obj'], temp_output, translation, reference_id, grid_id)
|
||
|
||
merged_obj = temp_output
|
||
|
||
# 最终结果
|
||
final_obj = os.path.join(output_model_dir, "merged_model.obj")
|
||
try:
|
||
if os.path.exists(final_obj):
|
||
os.remove(final_obj)
|
||
os.rename(merged_obj, final_obj)
|
||
except Exception as e:
|
||
self.logger.warning(f"重命名最终文件失败: {str(e)}")
|
||
shutil.copy2(merged_obj, final_obj)
|
||
try:
|
||
os.remove(merged_obj)
|
||
except:
|
||
pass
|
||
|
||
# 清理所有临时文件
|
||
for temp_file in temp_files:
|
||
if os.path.exists(temp_file):
|
||
try:
|
||
os.remove(temp_file)
|
||
except Exception as e:
|
||
self.logger.warning(
|
||
f"删除临时文件失败: {temp_file}, 错误: {str(e)}")
|
||
|
||
self.logger.info(
|
||
f"模型合并完成，输出目录: {output_model_dir}\n"
|
||
f"- OBJ文件: merged_model.obj\n"
|
||
f"- MTL文件: merged_model.mtl\n"
|
||
f"- 纹理文件: {len(os.listdir(output_model_dir)) - 2}个"
|
||
)
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"合并过程中发生错误: {str(e)}", exc_info=True)
|
||
raise
|