268 lines
10 KiB
Python
268 lines
10 KiB
Python
import os
|
||
import logging
|
||
import numpy as np
|
||
from typing import Dict
|
||
import pandas as pd
|
||
import shutil
|
||
|
||
|
||
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文件，返回顶点列表和面列表"""
|
||
vertices = []
|
||
faces = []
|
||
|
||
with open(file_path, 'r') as file:
|
||
for line in file:
|
||
parts = line.split()
|
||
if len(parts) == 0:
|
||
continue
|
||
if parts[0] == 'v': # 顶点
|
||
vertices.append(
|
||
[float(parts[1]), float(parts[2]), float(parts[3])])
|
||
elif parts[0] == 'f': # 面
|
||
faces.append([int(parts[1].split(
|
||
'/')[0]), int(parts[2].split('/')[0]), int(parts[3].split('/')[0])])
|
||
|
||
return vertices, faces
|
||
|
||
def write_obj(self, file_path, vertices, faces):
|
||
"""将修改后的顶点和面列表写入到.obj文件"""
|
||
with open(file_path, 'w') as file:
|
||
for vertex in vertices:
|
||
file.write(f"v {vertex[0]} {vertex[1]} {vertex[2]}\n")
|
||
|
||
for face in faces:
|
||
file.write(f"f {face[0]} {face[1]} {face[2]}\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):
|
||
"""合并两个OBJ文件"""
|
||
try:
|
||
self.logger.info(f"开始合并OBJ模型:\n输入1: {obj1_path}\n输入2: {obj2_path}")
|
||
|
||
# 检查输入文件是否存在
|
||
if not os.path.exists(obj1_path) or not os.path.exists(obj2_path):
|
||
raise FileNotFoundError("输入模型文件不存在")
|
||
|
||
# 读取两个obj文件
|
||
vertices1, faces1 = self.read_obj(obj1_path)
|
||
vertices2, faces2 = self.read_obj(obj2_path)
|
||
|
||
# 平移第二个模型的顶点
|
||
vertices2_translated = self.translate_vertices(
|
||
vertices2, translation)
|
||
|
||
# 合并顶点和面
|
||
all_vertices = vertices1 + vertices2_translated
|
||
all_faces = faces1 + \
|
||
[[f[0] + len(vertices1), f[1] + len(vertices1),
|
||
f[2] + len(vertices1)] for f in faces2]
|
||
|
||
# 写入合并后的obj文件
|
||
self.write_obj(output_path, all_vertices, all_faces)
|
||
|
||
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, file_path):
|
||
"""读取.mtl文件内容"""
|
||
with open(file_path, 'r') as file:
|
||
return file.read()
|
||
|
||
def copy_texture_files(self, src_dir: str, dst_dir: str, grid_id: tuple):
|
||
"""复制并重命名纹理文件
|
||
Args:
|
||
src_dir: 源纹理文件目录
|
||
dst_dir: 目标纹理文件目录
|
||
grid_id: 网格ID，用于重命名
|
||
"""
|
||
# 确保目标目录存在
|
||
os.makedirs(dst_dir, exist_ok=True)
|
||
|
||
# 复制所有png文件并重命名
|
||
for file in os.listdir(src_dir):
|
||
if file.endswith('.png'):
|
||
src_file = os.path.join(src_dir, file)
|
||
# 在文件名前添加网格ID前缀
|
||
new_name = f"grid_{grid_id[0]}_{grid_id[1]}_{file}"
|
||
dst_file = os.path.join(dst_dir, new_name)
|
||
shutil.copy2(src_file, dst_file)
|
||
self.logger.debug(f"复制纹理文件: {file} -> {new_name}")
|
||
|
||
return dst_dir
|
||
|
||
def update_mtl_content(self, mtl_content: str, grid_id: tuple) -> str:
|
||
"""更新MTL文件内容，修改纹理文件路径
|
||
Args:
|
||
mtl_content: 原MTL文件内容
|
||
grid_id: 网格ID，用于重命名纹理文件
|
||
Returns:
|
||
更新后的MTL文件内容
|
||
"""
|
||
updated_content = []
|
||
for line in mtl_content.split('\n'):
|
||
if line.startswith('map_Kd'): # 纹理文件路径行
|
||
# 获取原始文件名
|
||
original_file = line.split()[-1]
|
||
# 添加网格ID前缀
|
||
new_file = f"grid_{grid_id[0]}_{grid_id[1]}_{os.path.basename(original_file)}"
|
||
# 更新行内容
|
||
line = f"map_Kd {new_file}"
|
||
updated_content.append(line)
|
||
return '\n'.join(updated_content)
|
||
|
||
def merge_grid_obj(self, grid_points: Dict[tuple, pd.DataFrame], translations: Dict[tuple, tuple]):
|
||
"""合并所有网格的OBJ模型和纹理"""
|
||
self.logger.info("开始合并所有网格的OBJ模型")
|
||
|
||
if len(grid_points) < 2:
|
||
self.logger.info("只有一个网格，无需合并")
|
||
return
|
||
|
||
try:
|
||
# 创建输出目录
|
||
output_model_dir = os.path.join(self.output_dir, "merged_model")
|
||
os.makedirs(output_model_dir, exist_ok=True)
|
||
|
||
# 获取所有有效的网格OBJ文件
|
||
grid_objs = {}
|
||
for grid_id, points in grid_points.items():
|
||
grid_base_dir = os.path.join(
|
||
self.output_dir,
|
||
f"grid_{grid_id[0]}_{grid_id[1]}",
|
||
"project",
|
||
"odm_texturing"
|
||
)
|
||
grid_obj = os.path.join(grid_base_dir, "odm_textured_model_geo.obj")
|
||
grid_mtl = os.path.join(grid_base_dir, "odm_textured_model_geo.mtl")
|
||
|
||
if not os.path.exists(grid_obj) or not os.path.exists(grid_mtl):
|
||
self.logger.warning(
|
||
f"网格 ({grid_id[0]},{grid_id[1]}) 的OBJ或MTL文件不存在")
|
||
continue
|
||
|
||
grid_objs[grid_id] = {
|
||
'obj': grid_obj,
|
||
'mtl': grid_mtl,
|
||
'base_dir': grid_base_dir
|
||
}
|
||
|
||
if not grid_objs:
|
||
self.logger.error("没有找到有效的OBJ文件")
|
||
return
|
||
|
||
# 使用第一个网格作为参考
|
||
reference_id = list(grid_objs.keys())[0]
|
||
merged_obj = grid_objs[reference_id]['obj']
|
||
|
||
# 复制参考网格的纹理文件
|
||
self.copy_texture_files(
|
||
grid_objs[reference_id]['base_dir'],
|
||
output_model_dir,
|
||
reference_id
|
||
)
|
||
|
||
# 复制并更新参考网格的MTL文件
|
||
ref_mtl_content = self.read_mtl(grid_objs[reference_id]['mtl'])
|
||
updated_mtl = self.update_mtl_content(ref_mtl_content, reference_id)
|
||
|
||
self.logger.info(
|
||
f"使用网格 ({reference_id[0]},{reference_id[1]}) 作为参考网格")
|
||
|
||
# 依次合并其他网格
|
||
for grid_id, grid_files in list(grid_objs.items())[1:]:
|
||
# 复制当前网格的纹理文件
|
||
self.copy_texture_files(
|
||
grid_files['base_dir'],
|
||
output_model_dir,
|
||
grid_id
|
||
)
|
||
|
||
# 更新当前网格的MTL内容
|
||
current_mtl = self.read_mtl(grid_files['mtl'])
|
||
updated_mtl += '\n' + self.update_mtl_content(current_mtl, grid_id)
|
||
|
||
# 获取平移量
|
||
translation = translations[grid_id]
|
||
translation = (translation[0], translation[1], 0) # 添加z轴的0平移
|
||
|
||
output_obj = os.path.join(
|
||
output_model_dir,
|
||
f"merged_model_{reference_id[0]}_{reference_id[1]}_{grid_id[0]}_{grid_id[1]}.obj"
|
||
)
|
||
|
||
self.logger.info(
|
||
f"合并网格 ({grid_id[0]},{grid_id[1]}):\n"
|
||
f"平移量: x={translation[0]:.2f}m, y={translation[1]:.2f}m\n"
|
||
f"输出: {output_obj}"
|
||
)
|
||
|
||
self.merge_two_objs(merged_obj, grid_files['obj'],
|
||
output_obj, translation)
|
||
merged_obj = output_obj
|
||
|
||
# 最后的结果
|
||
final_obj = os.path.join(output_model_dir, "merged_model.obj")
|
||
final_mtl = os.path.join(output_model_dir, "merged_model.mtl")
|
||
|
||
# 保存最终的OBJ和MTL文件
|
||
if os.path.exists(merged_obj) and merged_obj != final_obj:
|
||
shutil.copy2(merged_obj, final_obj)
|
||
os.remove(merged_obj)
|
||
|
||
# 保存合并后的MTL文件
|
||
with open(final_mtl, 'w') as f:
|
||
f.write(updated_mtl)
|
||
|
||
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}个PNG文件"
|
||
)
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"OBJ模型合并过程中发生错误: {str(e)}", exc_info=True)
|
||
raise
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import sys
|
||
sys.path.append(os.path.dirname(
|
||
os.path.dirname(os.path.abspath(__file__))))
|
||
from utils.logger import setup_logger
|
||
import pandas as pd
|
||
|
||
# 设置输出目录和日志
|
||
output_dir = r"G:\ODM_output\1009"
|
||
setup_logger(output_dir)
|
||
|
||
# 构造测试用的grid_points字典
|
||
grid_points = {
|
||
(0, 0): pd.DataFrame({
|
||
'latitude': [39.9, 39.91],
|
||
'longitude': [116.3, 116.31],
|
||
'altitude': [100, 101]
|
||
}),
|
||
(0, 1): pd.DataFrame({
|
||
'latitude': [39.92, 39.93],
|
||
'longitude': [116.32, 116.33],
|
||
'altitude': [102, 103]
|
||
})
|
||
}
|
||
|
||
# 创建MergeObj实例并执行合并
|
||
merge_obj = MergeObj(output_dir)
|
||
merge_obj.merge_grid_obj(grid_points)
|