UAV/post_pro/conv_obj.py

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


class ConvertOBJ:
    def __init__(self, output_dir: str):
        self.output_dir = output_dir
        # 用于存储所有grid的UTM范围
        self.ref_east = float('inf')
        self.ref_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坐标作为参照系
        first_grid_id = list(grid_points.keys())[0]
        first_grid_dir = os.path.join(
            self.output_dir,
            f"grid_{first_grid_id[0]}_{first_grid_id[1]}",
            "project"
        )
        log_file = os.path.join(
            first_grid_dir, "odm_orthophoto", "odm_orthophoto_log.txt")
        self.ref_east, self.ref_north = self.read_utm_offset(log_file)

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

        self._create_merged_metadata()

    def _convert_single_grid(self, grid_id, grid_points):
        """转换单个网格的OBJ文件"""
        # 构建相关路径
        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")
        texturing_dst_dir = os.path.join(project_dir, "odm_texturing_dst")
        opensfm_dir = os.path.join(project_dir, "opensfm")
        log_file = os.path.join(
            project_dir, "odm_orthophoto", "odm_orthophoto_log.txt")
        os.makedirs(texturing_dst_dir, exist_ok=True)

        # 修改obj文件z坐标的值
        min_25d_z = self.get_min_z_from_obj(os.path.join(
            project_dir, 'odm_texturing_25d', 'odm_textured_model_geo.obj'))
        self.modify_z_in_obj(texturing_dir, min_25d_z)

        # 在新文件夹下，利用UTM偏移量，修改obj文件顶点坐标，纹理文件下采样
        utm_offset = self.read_utm_offset(log_file)
        modified_obj = self.modify_obj_coordinates(
            texturing_dir, texturing_dst_dir, utm_offset)
        self.downsample_texture(texturing_dir, texturing_dst_dir)

        # 执行格式转换，Linux下osgconv有问题，记得注释掉
        self.logger.info(f"开始转换网格 {grid_id} 的OBJ文件")
        output_osgb = os.path.join(texturing_dst_dir, "Tile.osgb")
        cmd = (
            f"osgconv {modified_obj} {output_osgb} "
            f"--compressed --smooth --fix-transparency "
        )
        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)}")

        # 创建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)
        target_osgb = os.path.join(
            tile_dir, f"Tile_{grid_id[0]}_{grid_id[1]}.osgb")
        shutil.copy2(output_osgb, target_osgb)

    def _create_merged_metadata(self):
        """创建合并后的metadata.xml文件"""
        # 转换为WGS84经纬度
        center_lon, center_lat = self.transformer.transform(
            self.ref_east, self.ref_north)
        metadata_content = f"""<?xml version="1.0" encoding="utf-8"?>
<ModelMetadata version="1">
    <SRS>EPSG:4326</SRS>
    <SRSOrigin>{center_lon},{center_lat},0</SRSOrigin>
        <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, texturing_dir: str, texturing_dst_dir: str, utm_offset: tuple) -> str:
        """修改obj文件中的顶点坐标，使用相对坐标系"""
        obj_file = os.path.join(
            texturing_dir, "odm_textured_model_modified.obj")
        obj_dst_file = os.path.join(
            texturing_dst_dir, "odm_textured_model_geo_utm.obj")
        if not os.path.exists(obj_file):
            raise FileNotFoundError(f"找不到OBJ文件: {obj_file}")
        shutil.copy2(os.path.join(texturing_dir, "odm_textured_model_geo.mtl"),
                     os.path.join(texturing_dst_dir, "odm_textured_model_geo.mtl"))
        east_offset, north_offset = utm_offset
        self.logger.info(
            f"UTM坐标偏移：{east_offset - self.ref_east}, {north_offset - self.ref_north}")

        try:
            with open(obj_file, 'r') as f_in, open(obj_dst_file, '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.ref_east)
                        y = float(parts[2]) + (north_offset - self.ref_north)
                        z = float(parts[3])
                        f_out.write(f'v {x:.6f} {z:.6f} {-y:.6f}\n')
                    elif line.startswith('vn '):  # 处理法线向量
                        parts = line.split()
                        nx = float(parts[1])
                        ny = float(parts[2])
                        nz = float(parts[3])
                        # 同步反转法线的 Y 轴
                        new_line = f"vn {nx} {nz} {-ny}\n"
                        f_out.write(new_line)
                    else:
                        # 其他行直接写入
                        f_out.write(line)

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

    def downsample_texture(self, src_dir: str, dst_dir: str):
        """复制并重命名纹理文件，对大于100MB的文件进行多次下采样，直到文件小于100MB
        Args:
            src_dir: 源纹理目录
            dst_dir: 目标纹理目录
        """
        for file in os.listdir(src_dir):
            if file.lower().endswith(('.png')):
                src_path = os.path.join(src_dir, file)
                dst_path = os.path.join(dst_dir, file)

                # 检查文件大小（以字节为单位）
                file_size = os.path.getsize(src_path)
                if file_size <= 100 * 1024 * 1024:  # 如果文件小于等于100MB，直接复制
                    shutil.copy2(src_path, dst_path)
                else:
                    # 文件大于100MB，进行下采样
                    img = cv2.imread(src_path, cv2.IMREAD_UNCHANGED)
                    if_first_ds = True
                    while file_size > 100 * 1024 * 1024:  # 大于100MB
                        self.logger.info(f"纹理文件 {file} 大于100MB，进行下采样")

                        if if_first_ds:
                            # 计算新的尺寸（长宽各变为1/4）
                            new_size = (img.shape[1] // 4,
                                        img.shape[0] // 4)  # 逐步减小尺寸
                            # 使用双三次插值进行下采样
                            resized_img = cv2.resize(
                                img, new_size, interpolation=cv2.INTER_CUBIC)
                            if_first_ds = False
                        else:
                            # 计算新的尺寸（长宽各变为1/2）
                            new_size = (img.shape[1] // 2,
                                        img.shape[0] // 2)  # 逐步减小尺寸
                            # 使用双三次插值进行下采样
                            resized_img = cv2.resize(
                                img, new_size, interpolation=cv2.INTER_CUBIC)

                        # 更新文件路径为下采样后的路径
                        cv2.imwrite(dst_path, resized_img, [
                                    cv2.IMWRITE_PNG_COMPRESSION, 9])

                        # 更新文件大小和图像
                        file_size = os.path.getsize(dst_path)
                        img = cv2.imread(dst_path, cv2.IMREAD_UNCHANGED)
                        self.logger.info(
                            f"下采样后文件大小: {file_size / (1024 * 1024):.2f} MB")

    def get_min_z_from_obj(self, file_path):
        min_z = float('inf')  # 初始值设为无穷大
        with open(file_path, 'r') as obj_file:
            for line in obj_file:
                # 检查每一行是否是顶点定义（以 'v ' 开头）
                if line.startswith('v '):
                    # 获取顶点坐标
                    parts = line.split()
                    # 将z值转换为浮动数字
                    z = float(parts[3])
                    # 更新最小z值
                    if z < min_z:
                        min_z = z
        return min_z

    def modify_z_in_obj(self, texturing_dir, min_25d_z):
        obj_file = os.path.join(texturing_dir, 'odm_textured_model_geo.obj')
        output_file = os.path.join(
            texturing_dir, 'odm_textured_model_modified.obj')
        with open(obj_file, 'r') as f_in, open(output_file, 'w') as f_out:
            for line in f_in:
                if line.startswith('v '):  # 顶点坐标行
                    parts = line.strip().split()
                    x = float(parts[1])
                    y = float(parts[2])
                    z = float(parts[3])

                    if z < min_25d_z:
                        z = min_25d_z

                    f_out.write(f"v {x} {y} {z}\n")
                else:
                    f_out.write(line)