UAV/utils/odm_monitor.py

import os
import logging
import docker
from typing import Tuple


class ODMProcessMonitor:
    """ODM进程监控器"""

    def __init__(self, max_retries: int = 3, mode: str = "快拼模式"):
        self.max_retries = max_retries
        self.logger = logging.getLogger('UAV_Preprocess.ODMMonitor')
        self.mode = mode
        self.client = docker.from_env()

    def _check_success(self, grid_dir: str) -> bool:
        """检查ODM是否执行成功"""
        success_markers = ['odm_orthophoto', 'odm_georeferencing']
        if self.mode != "快拼模式":
            success_markers.append('odm_texturing')
        return all(os.path.exists(os.path.join(grid_dir, 'project', marker)) for marker in success_markers)

    def run_odm_with_monitor(self, grid_dir: str, grid_idx: int, fast_mode: bool = True) -> Tuple[bool, str]:
        """运行ODM容器"""
        try:
            self.logger.info(f"开始处理网格 {grid_idx + 1}")

            # 准备容器配置
            volumes = {
                grid_dir: {'bind': '/datasets', 'mode': 'rw'}
            }

            # 准备命令参数
            command = [
                "--project-path", "/datasets", "project",
                "--max-concurrency", "10",
                "--force-gps",
                "--rerun-all"
            ]

            if fast_mode:
                command.extend([
                    "--feature-quality", "lowest",
                    "--orthophoto-resolution", "8",
                    "--fast-orthophoto",
                    "--skip-3dmodel"
                ])

            # 运行容器并等待完成
            container = self.client.containers.run(
                "opendronemap/odm",
                command=command,
                volumes=volumes,
                detach=True,
                remove=True,
                environment={"PYTHONUNBUFFERED": "1"},
                mem_limit="0",  # 不限制内存
                cpu_count=0,    # 使用所有CPU
                network_mode="host"  # 使用主机网络模式
            )

            # 等待容器完成并获取状态码
            result = container.wait()

            # 检查是否成功完成
            if result['StatusCode'] == 0 and self._check_success(grid_dir):
                self.logger.info(f"网格 {grid_idx + 1} 处理成功")
                return True, ""

            return False, f"网格 {grid_idx + 1} 处理失败"

        except Exception as e:
            error_msg = f"执行异常: {str(e)}"
            self.logger.error(error_msg)
            return False, error_msg