UAV/post_pro/merge_tif.py

from osgeo import gdal
import logging
import os
from typing import Dict
import pandas as pd
import time
import shutil


class MergeTif:
    def __init__(self, output_dir: str):
        self.output_dir = output_dir
        self.logger = logging.getLogger('UAV_Preprocess.MergeTif')

    def merge_two_tifs(self, input_tif1: str, input_tif2: str, output_tif: str):
        """合并两张TIF影像"""
        try:
            self.logger.info("开始合并TIF影像")
            self.logger.info(f"输入影像1: {input_tif1}")
            self.logger.info(f"输入影像2: {input_tif2}")
            self.logger.info(f"输出影像: {output_tif}")

            # 检查输入文件是否存在
            if not os.path.exists(input_tif1) or not os.path.exists(input_tif2):
                error_msg = "输入影像文件不存在"
                self.logger.error(error_msg)
                raise FileNotFoundError(error_msg)

            # 打开影像，检查投影是否一致
            datasets = []
            try:
                for tif in [input_tif1, input_tif2]:
                    ds = gdal.Open(tif)
                    if ds is None:
                        error_msg = f"无法打开影像文件: {tif}"
                        self.logger.error(error_msg)
                        raise ValueError(error_msg)
                    datasets.append(ds)

                projections = [ds.GetProjection() for ds in datasets]
                self.logger.debug(f"影像1投影: {projections[0]}")
                self.logger.debug(f"影像2投影: {projections[1]}")

                # 检查投影是否一致
                if len(set(projections)) != 1:
                    error_msg = "影像的投影不一致，请先进行重投影！"
                    self.logger.error(error_msg)
                    raise ValueError(error_msg)

                # 如果输出文件已存在，先删除
                if os.path.exists(output_tif):
                    try:
                        os.remove(output_tif)
                    except Exception as e:
                        self.logger.warning(f"删除已存在的输出文件失败: {str(e)}")
                        # 生成一个新的输出文件名
                        base, ext = os.path.splitext(output_tif)
                        output_tif = f"{base}_{int(time.time())}{ext}"
                        self.logger.info(f"使用新的输出文件名: {output_tif}")

                # 创建 GDAL Warp 选项
                warp_options = gdal.WarpOptions(
                    format="GTiff",
                    resampleAlg="average",
                    srcNodata=0,
                    dstNodata=0,
                    multithread=True
                )

                self.logger.info("开始执行影像拼接...")
                result = gdal.Warp(output_tif, datasets, options=warp_options)

                if result is None:
                    error_msg = "影像拼接失败"
                    self.logger.error(error_msg)
                    raise RuntimeError(error_msg)

                # 获取输出影像的基本信息
                output_dataset = gdal.Open(output_tif)
                if output_dataset:
                    width = output_dataset.RasterXSize
                    height = output_dataset.RasterYSize
                    bands = output_dataset.RasterCount
                    self.logger.info(
                        f"拼接完成，输出影像大小: {width}x{height}，波段数: {bands}")
                    output_dataset = None  # 显式关闭数据集

                self.logger.info(f"影像拼接成功，输出文件保存至: {output_tif}")

            finally:
                # 确保所有数据集都被正确关闭
                for ds in datasets:
                    if ds:
                        ds = None
                result = None

        except Exception as e:
            self.logger.error(f"影像拼接过程中发生错误: {str(e)}", exc_info=True)
            raise

    def merge_grid_tif(self, grid_points: Dict[tuple, pd.DataFrame], product_info: dict):
        """合并指定产品的所有网格"""
        product_name = product_info['name']
        product_path = product_info['path']
        filename = product_info['filename']

        self.logger.info(f"开始合并{product_name}")

        if len(grid_points) < 2:
            self.logger.info("只有一个网格，无需合并")
            return

        input_tif1, input_tif2 = None, None
        merge_count = 0

        try:
            for grid_id, points in grid_points.items():
                grid_tif = os.path.join(
                    self.output_dir,
                    f"grid_{grid_id[0]}_{grid_id[1]}",
                    "project",
                    product_path,
                    filename
                )

                if not os.path.exists(grid_tif):
                    self.logger.warning(
                        f"网格 ({grid_id[0]},{grid_id[1]}) 的{product_name}不存在: {grid_tif}")
                    continue

                # 如果文件大于600MB，则不使用original文件
                file_size_mb = os.path.getsize(grid_tif) / (1024 * 1024)  # 转换为MB
                if file_size_mb > 600:
                    grid_tif = os.path.join(
                        self.output_dir,
                        f"grid_{grid_id[0]}_{grid_id[1]}",
                        "project",
                        product_path,
                        "odm_orthophoto.tif"
                    )

                if input_tif1 is None:
                    input_tif1 = grid_tif
                    self.logger.info(f"设置第一个输入{product_name}: {input_tif1}")
                else:
                    input_tif2 = grid_tif
                    # 生成带时间戳的临时输出文件名
                    temp_output = os.path.join(
                        self.output_dir,
                        f"temp_merged_{int(time.time())}_{product_info['output']}"
                    )

                    self.logger.info(
                        f"开始合并{product_name}第 {merge_count + 1} 次:\n"
                        f"输入1: {input_tif1}\n"
                        f"输入2: {input_tif2}\n"
                        f"输出: {temp_output}"
                    )

                    self.merge_two_tifs(input_tif1, input_tif2, temp_output)
                    merge_count += 1

                    # 如果之前的输入文件是临时文件，则删除它
                    if 'temp_merged_' in input_tif1:
                        try:
                            os.remove(input_tif1)
                        except Exception as e:
                            self.logger.warning(f"删除临时文件失败: {str(e)}")

                    input_tif1 = temp_output
                    input_tif2 = None

            # 重命名最终的临时文件为目标文件名
            final_output = os.path.join(
                self.output_dir, product_info['output'])
            if os.path.exists(final_output):
                try:
                    os.remove(final_output)
                except Exception as e:
                    self.logger.warning(f"删除已存在的最终输出文件失败: {str(e)}")
                    final_output = os.path.join(
                        self.output_dir,
                        f"merged_{int(time.time())}_{product_info['output']}"
                    )

            try:
                os.rename(input_tif1, final_output)
            except Exception as e:
                self.logger.warning(f"重命名最终文件失败: {str(e)}")
                shutil.copy2(input_tif1, final_output)
                try:
                    os.remove(input_tif1)
                except:
                    pass

            self.logger.info(
                f"{product_name}合并完成，共执行 {merge_count} 次合并，"
                f"最终输出文件: {final_output}"
            )

        except Exception as e:
            self.logger.error(
                f"{product_name}合并过程中发生错误: {str(e)}", exc_info=True)
            raise

    def merge_all_tifs(self, grid_points: Dict[tuple, pd.DataFrame], produce_dem: bool):
        """合并所有产品（正射影像、DSM和DTM）"""
        try:
            products = [
                {
                    'name': '正射影像',
                    'path': 'odm_orthophoto',
                    'filename': 'odm_orthophoto.original.tif',
                    'output': 'orthophoto.tif'
                },

            ]
            if produce_dem:
                products.append(
                    {
                        'name': 'DSM',
                        'path': 'odm_dem',
                        'filename': 'dsm.original.tif',
                        'output': 'dsm.tif'
                    },
                    {
                        'name': 'DTM',
                        'path': 'odm_dem',
                        'filename': 'dtm.original.tif',
                        'output': 'dtm.tif'
                    }
                )

            for product in products:
                self.merge_grid_tif(grid_points, product)

            self.logger.info("所有产品合并完成")
        except Exception as e:
            self.logger.error(f"产品合并过程中发生错误: {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字典
    # 假设我们有两个网格，每个网格包含一些GPS点的DataFrame
    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]
        })
    }

    # 创建MergeTif实例并执行合并
    merge_tif = MergeTif(output_dir)
    merge_tif.merge_all_tifs(grid_points)