UAV/post_pro/obj_post_pro.py

import os
import logging
import subprocess
import shutil
from typing import List, Tuple
import numpy as np


class ObjPostProcessor:
    def __init__(self, output_dir: str):
        self.output_dir = output_dir
        self.logger = logging.getLogger('UAV_Preprocess.ObjPostProcessor')
        # 用于存储所有grid的UTM范围
        self.min_east = float('inf')
        self.min_north = float('inf')
        self.max_east = float('-inf')
        self.max_north = float('-inf')

    def process_obj_files(self):
        """处理所有grid中的obj文件"""
        try:
            # 1. 遍历所有grid文件夹
            grid_dirs = [d for d in os.listdir(
                self.output_dir) if d.startswith('grid_')]

            # 第一次遍历：获取所有grid的UTM范围
            for grid_dir in grid_dirs:
                grid_path = os.path.join(self.output_dir, grid_dir)
                log_file = os.path.join(
                    grid_path, 'project', 'odm_orthophoto', 'odm_orthophoto_log.txt')
                east_offset, north_offset = self.read_utm_offset(log_file)
                
                # 更新UTM范围
                self.min_east = min(self.min_east, east_offset)
                self.min_north = min(self.min_north, north_offset)
                self.max_east = max(self.max_east, east_offset)
                self.max_north = max(self.max_north, north_offset)

            # 创建osgb输出目录
            osgb_dir = os.path.join(self.output_dir, 'osgb')
            os.makedirs(os.path.join(osgb_dir, 'Data'), exist_ok=True)

            # 第二次遍历：处理每个grid
            for grid_dir in grid_dirs:
                grid_path = os.path.join(self.output_dir, grid_dir)
                self.process_single_grid(grid_path, osgb_dir)

            # 创建metadata.xml
            self.create_metadata_xml(osgb_dir)

            self.logger.info("所有grid处理完成")
            return True
        except Exception as e:
            self.logger.error(f"处理obj文件时发生错误: {str(e)}")
            return False

    def process_single_grid(self, grid_path: str, osgb_dir: str):
        """处理单个grid的obj文件"""
        try:
            # 1. 读取UTM偏移量
            log_file = os.path.join(
                grid_path, 'project', 'odm_orthophoto', 'odm_orthophoto_log.txt')
            utm_offset = self.read_utm_offset(log_file)

            # 2. 修改obj文件的顶点坐标
            obj_file = os.path.join(
                grid_path, 'project', 'odm_texturing', 'odm_textured_model_geo.obj')
            modified_obj = self.modify_obj_coordinates(obj_file, utm_offset)

            # 3. 使用osgconv转换为osgb
            grid_name = os.path.basename(grid_path)
            self.convert_to_osgb(modified_obj, grid_name, osgb_dir)

        except Exception as e:
            self.logger.error(f"处理grid {grid_path} 时发生错误: {str(e)}")
            raise

    def read_utm_offset(self, log_file: str) -> Tuple[float, float]:
        """读取UTM偏移量"""
        try:
            east_offset = None
            north_offset = None

            with open(log_file, 'r') as f:
                lines = f.readlines()
                for i, line in enumerate(lines):
                    if 'utm_north_offset' in line and i + 1 < len(lines):
                        north_offset = float(lines[i + 1].strip())
                    elif 'utm_east_offset' in line and i + 1 < len(lines):
                        east_offset = float(lines[i + 1].strip())

            if east_offset is None or north_offset is None:
                raise ValueError("未找到UTM偏移量")

            return east_offset, north_offset
        except Exception as e:
            self.logger.error(f"读取UTM偏移量时发生错误: {str(e)}")
            raise

    def modify_obj_coordinates(self, obj_file: str, utm_offset: Tuple[float, float]) -> str:
        """修改obj文件中的顶点坐标，使用相对坐标系"""
        east_offset, north_offset = utm_offset
        output_obj = obj_file.replace('.obj', '_utm.obj')

        try:
            with open(obj_file, 'r') as f_in, open(output_obj, 'w') as f_out:
                for line in f_in:
                    if line.startswith('v '):
                        # 处理顶点坐标行
                        parts = line.strip().split()
                        # 使用相对于整体最小UTM坐标的偏移
                        x = float(parts[1]) + (east_offset - self.min_east)
                        y = float(parts[2]) + (north_offset - self.min_north)
                        z = float(parts[3])
                        f_out.write(f'v {x:.6f} {y:.6f} {z:.6f}\n')
                    else:
                        # 其他行直接写入
                        f_out.write(line)

            # 复制材质文件
            mtl_file = obj_file.replace('.obj', '.mtl')
            if os.path.exists(mtl_file):
                shutil.copy2(mtl_file, mtl_file.replace('.mtl', '_utm.mtl'))

            return output_obj
        except Exception as e:
            self.logger.error(f"修改obj坐标时发生错误: {str(e)}")
            raise

    def convert_to_osgb(self, obj_file: str, grid_name: str, osgb_dir: str):
        """使用osgconv将obj转换为osgb"""
        try:
            # 创建tile目录
            tile_dir = os.path.join(osgb_dir, 'Data', grid_name)
            os.makedirs(tile_dir, exist_ok=True)

            output_osgb = os.path.join(tile_dir, f'{grid_name}.osgb')

            # 构建osgconv命令
            cmd = [
                'osgconv',
                '--compressed',
                '--smooth',
                '--fix-transparency',
                '-o', '0,1,0-0,0,-1',
                obj_file,
                output_osgb
            ]

            # 执行命令
            self.logger.info(f"执行osgconv命令：{cmd}")
            result = subprocess.run(cmd, capture_output=True, text=True)

            if result.returncode != 0:
                raise Exception(f"osgconv执行失败: {result.stderr}")

            self.logger.info(f"转换完成: {output_osgb}")

        except Exception as e:
            self.logger.error(f"转换osgb时发生错误: {str(e)}")
            raise

    def create_metadata_xml(self, osgb_dir: str):
        """创建metadata.xml文件，包含UTM偏移信息"""
        try:
            # 这里需要将UTM坐标转换为WGS84经纬度坐标
            # 这里使用示例值，实际应用中需要进行真实的坐标转换
            metadata_content = f'''<?xml version="1.0" encoding="utf-8"?>
<ModelMetadata version="1">
    <SRS>EPSG:32649</SRS>
    <SRSOrigin>{self.min_east:.6f},{self.min_north:.6f},0.000000</SRSOrigin>
    <Texture>
        <ColorSource>Visible</ColorSource>
    </Texture>
</ModelMetadata>'''
            
            with open(os.path.join(osgb_dir, 'metadata.xml'), 'w') as f:
                f.write(metadata_content)
                
        except Exception as e:
            self.logger.error(f"创建metadata.xml时发生错误: {str(e)}")
            raise