UAV/post_pro/merge_obj.py

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)