UAV/post_pro/conv_obj.py

import os
import subprocess
import json
import shutil
import logging
from pyproj import Transformer


class ConvertOBJ:
    def __init__(self, output_dir: str, center_lat: float, center_lon: float):
        self.output_dir = output_dir
        self.center_lat = center_lat
        self.center_lon = center_lon
        # 用于存储所有grid的UTM范围
        self.min_east = float('inf')
        self.min_north = float('inf')
        self.max_east = float('-inf')
        self.max_north = float('-inf')
        # 初始化UTM到WGS84的转换器
        self.transformer = Transformer.from_crs(
            "EPSG:32649", "EPSG:4326", always_xy=True)
        self.logger = logging.getLogger('UAV_Preprocess.ConvertOBJ')

    def convert_grid_obj(self, grid_points):
        """转换每个网格的OBJ文件为OSGB格式"""
        os.makedirs(os.path.join(self.output_dir,
                    "osgb", "Data"), exist_ok=True)

        # 第一次遍历：获取所有grid的UTM范围
        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"
            )
            log_file = os.path.join(
                base_dir, "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)

        tile_infos = []
        for grid_id in grid_points.keys():
            try:
                tile_info = self._convert_single_grid(grid_id, grid_points)
                tile_infos.append(tile_info)
            except Exception as e:
                self.logger.error(f"网格 {grid_id} 转换失败: {str(e)}")

        self._create_merged_metadata(tile_infos)

    def _convert_single_grid(self, grid_id, grid_points):
        """转换单个网格的OBJ文件"""
        # 1. 构建相关路径
        grid_name = f"grid_{grid_id[0]}_{grid_id[1]}"
        project_dir = os.path.join(self.output_dir, grid_name, "project")
        texturing_dir = os.path.join(project_dir, "odm_texturing")
        opensfm_dir = os.path.join(project_dir, "opensfm")
        obj_file = os.path.join(texturing_dir, "odm_textured_model_geo.obj")
        log_file = os.path.join(
            project_dir, "odm_orthophoto", "odm_orthophoto_log.txt")
        if not os.path.exists(obj_file):
            raise FileNotFoundError(f"找不到OBJ文件: {obj_file}")

        # 2. 读取UTM偏移量，修改obj文件顶点坐标
        utm_offset = self.read_utm_offset(log_file)
        modified_obj = self.modify_obj_coordinates(
            obj_file, utm_offset)

        # 3. 执行格式转换
        self.logger.info(f"开始转换网格 {grid_id} 的OBJ文件")
        output_osgb = os.path.join(texturing_dir, "Tile.osgb")
        cmd = (
            f"osgconv {obj_file} {output_osgb} "
            f"--compressed --smooth --fix-transparency "
            f"-o 0,1,0-0,0,-1"
        )
        self.logger.info(f"执行osgconv命令：{cmd}")

        try:
            subprocess.run(cmd, shell=True, check=True, cwd=texturing_dir)
        except subprocess.CalledProcessError as e:
            raise RuntimeError(f"OSGB转换失败: {str(e)}")

        # 4. 读取地理信息
        ref_lla_file = os.path.join(opensfm_dir, "reference_lla.json")
        with open(ref_lla_file, 'r') as f:
            ref_lla = json.load(f)

        # 5. 创建OSGB目录结构
        osgb_base_dir = os.path.join(self.output_dir, "osgb")
        data_dir = os.path.join(osgb_base_dir, "Data")
        tile_dir = os.path.join(data_dir, f"Tile_{grid_id[0]}_{grid_id[1]}")
        os.makedirs(tile_dir, exist_ok=True)

        # 5. 复制OSGB文件
        target_osgb = os.path.join(
            tile_dir, f"Tile_{grid_id[0]}_{grid_id[1]}.osgb")
        shutil.copy2(output_osgb, target_osgb)

        # 计算当前网格的边界框
        grid_data = grid_points[grid_id]
        # 假设grid_data是一个列表，每个元素都是包含lat和lon的字典
        lats = [point['lat'] for point in grid_data]
        lons = [point['lon'] for point in grid_data]

        min_lat = min(lats)
        max_lat = max(lats)
        min_lon = min(lons)
        max_lon = max(lons)

        tile_info = {
            'id': f"{grid_id[0]}_{grid_id[1]}",
            'bounds': {
                'min_lat': min_lat,
                'max_lat': max_lat,
                'min_lon': min_lon,
                'max_lon': max_lon
            },
        }
        return tile_info

    def _calculate_distance(self, lat1, lon1, lat2, lon2):
        """计算两点间的距离（米）"""
        from math import sin, cos, sqrt, atan2, radians
        R = 6371000  # 地球半径（米）

        lat1, lon1, lat2, lon2 = map(radians, [lat1, lon1, lat2, lon2])
        dlat = lat2 - lat1
        dlon = lon2 - lon1

        a = sin(dlat/2)**2 + cos(lat1) * cos(lat2) * sin(dlon/2)**2
        c = 2 * atan2(sqrt(a), sqrt(1-a))
        return R * c

    def _create_merged_metadata(self, tile_infos):
        """创建合并后的metadata.xml文件"""
        metadata_content = f"""<?xml version="1.0" encoding="utf-8"?>
<ModelMetadata version="1">
    <SRS>EPSG:4326</SRS>
    <SRSOrigin>{self.center_lon},{self.center_lat},0.000000</SRSOrigin>
    <TileStructure>
        <RootNode>
            <BoundingBox>
                <MinLat>{min([t['bounds']['min_lat'] for t in tile_infos])}</MinLat>
                <MaxLat>{max([t['bounds']['max_lat'] for t in tile_infos])}</MaxLat>
                <MinLon>{min([t['bounds']['min_lon'] for t in tile_infos])}</MinLon>
                <MaxLon>{max([t['bounds']['max_lon'] for t in tile_infos])}</MaxLon>
            </BoundingBox>
            <Tiles>"""

        for tile in tile_infos:
            metadata_content += f"""
                    <Tile id="{tile['id']}">
                        <BoundingBox>
                            <MinLat>{tile['bounds']['min_lat']}</MinLat>
                            <MaxLat>{tile['bounds']['max_lat']}</MaxLat>
                            <MinLon>{tile['bounds']['min_lon']}</MinLon>
                            <MaxLon>{tile['bounds']['max_lon']}</MaxLon>
                        </BoundingBox>
                    </Tile>"""

        metadata_content += """
                </Tiles>
            </RootNode>
        </TileStructure>
        <Texture>
            <ColorSource>Visible</ColorSource>
        </Texture>
    </ModelMetadata>"""

        metadata_file = os.path.join(self.output_dir, "osgb", "metadata.xml")
        with open(metadata_file, 'w', encoding='utf-8') as f:
            f.write(metadata_content)

    def read_utm_offset(self, log_file: str) -> tuple:
        """读取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) -> 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)

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