UAV2/main.py

import os
import shutil
from datetime import timedelta
from dataclasses import dataclass
from typing import Dict, Tuple
import psutil
import pandas as pd
from pathlib import Path

from filter.cluster_filter import GPSCluster
from utils.gps_extractor import GPSExtractor
from utils.grid_divider import GridDivider
from utils.logger import setup_logger
from utils.visualizer import FilterVisualizer
from utils.docker_runner import DockerRunner
from post_pro.conv_obj import ConvertOBJ


@dataclass
class ProcessConfig:
    """配置类"""

    image_dir: str
    output_dir: str
    # 聚类过滤参数
    cluster_eps: float = 0.01
    cluster_min_samples: int = 5
    # 时间组重叠过滤参数
    time_group_overlap_threshold: float = 0.7
    time_group_interval: timedelta = timedelta(minutes=5)
    # 孤立点过滤参数
    filter_distance_threshold: float = 0.001  # 经纬度距离
    filter_min_neighbors: int = 6
    # 密集点过滤参数
    filter_grid_size: float = 0.001
    filter_dense_distance_threshold: float = 10  # 普通距离，单位：米
    filter_time_threshold: timedelta = timedelta(minutes=5)
    # 网格划分参数
    grid_overlap: float = 0.05
    grid_size: float = 500
    # 几个pipline过程是否开启
    mode: str = "快拼模式"
    accuracy: str = "medium"
    produce_dem: bool = False


class ODM_Plugin:
    def __init__(self, config: ProcessConfig):
        self.config = config
        self.logger = setup_logger(config.output_dir)

        # 检查磁盘空间
        # TODO 现在输入目录的磁盘空间也需要检查
        self._check_disk_space()

        # 清理并重建输出目录
        if os.path.exists(config.output_dir):
            self._clean_output_dir()
        self._setup_output_dirs()

        # 修改输入目录，符合ODM要求，从这里开始，image_dir就是project_path
        self._rename_input_dir()
        self.project_path = self.config.image_dir

        # 初始化其他组件
        self.gps_points = None
        self.grid_points = None
        self.visualizer = FilterVisualizer(config.output_dir)

    def _clean_output_dir(self):
        """清理输出目录"""
        try:
            shutil.rmtree(self.config.output_dir)
            self.logger.info(f"已清理输出目录: {self.config.output_dir}")
        except Exception as e:
            self.logger.info(f"清理输出目录时发生错误: {str(e)}")
            raise

    def _setup_output_dirs(self):
        """创建必要的输出目录结构"""
        try:
            # 创建主输出目录
            os.makedirs(self.config.output_dir)

            # 创建过滤图像保存目录
            os.makedirs(os.path.join(self.config.output_dir, 'filter_imgs'))

            # 创建日志目录
            os.makedirs(os.path.join(self.config.output_dir, 'logs'))

            self.logger.info(f"已创建输出目录结构: {self.config.output_dir}")
        except Exception as e:
            self.logger.info(f"创建输出目录时发生错误: {str(e)}")
            raise

    def _get_directory_size(self, path):
        """获取目录的总大小（字节）"""
        total_size = 0
        for dirpath, dirnames, filenames in os.walk(path):
            for filename in filenames:
                file_path = os.path.join(dirpath, filename)
                try:
                    total_size += os.path.getsize(file_path)
                except (OSError, FileNotFoundError):
                    continue
        return total_size

    def _check_disk_space(self):
        """检查磁盘空间是否足够"""
        # 获取输入目录大小
        input_size = self._get_directory_size(self.config.image_dir)

        # 获取输出目录所在磁盘的剩余空间
        output_drive = os.path.splitdrive(
            os.path.abspath(self.config.output_dir))[0]
        if not output_drive:  # 处理Linux/Unix路径
            output_drive = '/home'

        disk_usage = psutil.disk_usage(output_drive)
        free_space = disk_usage.free

        # 计算所需空间（输入大小的1.5倍）
        required_space = input_size * 12

        if free_space < required_space:
            error_msg = (
                f"磁盘空间不足!\n"
                f"输入目录大小: {input_size / (1024**3):.2f} GB\n"
                f"所需空间: {required_space / (1024**3):.2f} GB\n"
                f"可用空间: {free_space / (1024**3):.2f} GB\n"
                f"在驱动器 {output_drive}"
            )
            raise RuntimeError(error_msg)

    def _rename_input_dir(self):
        image_dir = Path(self.config.image_dir).resolve()

        if not image_dir.exists() or not image_dir.is_dir():
            raise ValueError(
                f"Provided path '{image_dir}' is not a valid directory.")

        # 原目录名和父路径
        parent_dir = image_dir.parent
        original_name = image_dir.name

        # 新的 images 路径（原目录重命名为 images）
        images_path = parent_dir / "images"

        # 重命名原目录为 images
        image_dir.rename(images_path)

        # 创建一个新的、和原目录同名的文件夹
        new_root = parent_dir / original_name
        new_root.mkdir(exist_ok=False)

        # 创建 project 子文件夹
        project_dir = new_root / "project"
        project_dir.mkdir()

        # 把 images 文件夹移动到 project 下
        final_images_path = project_dir / "images"
        shutil.move(str(images_path), str(final_images_path))

        self.logger.info(f"符合标准输入的文件夹结构已经创建好了，{final_images_path}")

        return final_images_path

    def extract_gps(self) -> pd.DataFrame:
        """提取GPS数据"""
        self.logger.info("开始提取GPS数据")
        extractor = GPSExtractor(self.project_path)
        self.gps_points = extractor.extract_all_gps()
        self.logger.info(f"成功提取 {len(self.gps_points)} 个GPS点")

    def cluster(self):
        """使用DBSCAN对GPS点进行聚类，只保留最大的类"""
        previous_points = self.gps_points.copy()
        clusterer = GPSCluster(
            self.gps_points,
            eps=self.config.cluster_eps,
            min_samples=self.config.cluster_min_samples
        )

        self.clustered_points = clusterer.fit()
        self.gps_points = clusterer.get_cluster_stats(self.clustered_points)

        self.visualizer.visualize_filter_step(
            self.gps_points, previous_points, "1-Clustering")

    def divide_grids(self):
        """划分网格
        Returns:
            tuple: (grid_points, translations)
                - grid_points: 网格点数据字典
                - translations: 网格平移量字典
        """
        grid_divider = GridDivider(
            overlap=self.config.grid_overlap,
            grid_size=self.config.grid_size,
            project_path=self.project_path,
            output_dir=self.config.output_dir
        )
        grids, self.grid_points = grid_divider.adjust_grid_size_and_overlap(
            self.gps_points
        )
        grid_divider.visualize_grids(self.gps_points, grids)
        grid_divider.save_image_groups(self.grid_points)
        if len(grids) >= 20:
            self.logger.warning("网格数量已超过20, 需要人工调整分区")

    def odm_docker_runner(self):
        """"运行OMD docker容器"""
        self.logger.info("开始运行Docker容器")
        # TODO：加一些容错处理
        docker_runner = DockerRunner(self.project_path)
        docker_runner.run_odm_container()

    def convert_obj(self):
        """转换OBJ模型"""
        self.logger.info("开始转换OBJ模型")
        converter = ConvertOBJ(self.config.output_dir)
        converter.convert_grid_obj(self.grid_points)

    def post_process(self):
        """后处理：合并或复制处理结果"""
        self.logger.info("开始后处理")

        self.logger.info("拷贝正射影像至输出目录")
        orthophoto_tif_path = os.path.join(
            self.project_path, "project", "odm_orthophoto", "odm_orthophoto.tif")
        shutil.copy(orthophoto_tif_path, self.config.output_dir)
        # if self.config.mode == "三维模式":
        #     self.convert_obj()
        # else:
        #     pass

    def process(self):
        """执行完整的预处理流程"""
        try:
            self.extract_gps()
            self.cluster()
            self.divide_grids()
            self.logger.info("==========预处理任务完成==========")
            self.odm_docker_runner()
            self.post_process()

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