UAV/utils/odm_monitor.py

import os
import time
import psutil
import logging
import subprocess
from typing import Optional, Tuple
from collections import deque


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

    def __init__(self, max_retries: int = 3, check_interval: int = 10, mode: str = "快拼模式"):
        """
        初始化监控器

        Args:
            max_retries: 最大重试次数
            check_interval: 检查间隔（秒）
            mode: 模式
        """
        self.max_retries = max_retries
        self.check_interval = check_interval
        self.logger = logging.getLogger('UAV_Preprocess.ODMMonitor')
        self.mode = mode
        # 用于存储最后N行输出
        self.last_outputs = deque(maxlen=50)

    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, command: str, grid_dir: str, grid_idx: int) -> Tuple[bool, str]:
        """运行ODM命令并监控进程"""
        attempt = 0
        while attempt < self.max_retries:
            try:
                self.logger.info(f"网格 {grid_idx + 1} 第 {attempt + 1} 次尝试执行ODM")

                # 使用 subprocess.Popen 启动进程
                process = subprocess.Popen(
                    command,
                    shell=True,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    bufsize=-1,  # 使用系统默认缓冲
                    text=True
                )

                # 等待进程完成并获取所有输出
                stdout, stderr = process.communicate()
                return_code = process.returncode

                # 只保存最后N行输出
                last_lines = stdout.strip().split('\n')[-50:]
                if last_lines:
                    self.logger.info(f"网格 {grid_idx + 1} 最后的输出:")
                    for line in last_lines:
                        self.logger.info(line)

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

                self.logger.warning(f"网格 {grid_idx + 1} 第 {attempt + 1} 次尝试失败")
                if stderr:
                    self.logger.error(f"错误输出: {stderr}")

            except Exception as e:
                error_msg = f"监控进程发生异常: {str(e)}"
                self.logger.error(error_msg)
                return False, error_msg

            attempt += 1
            if attempt < self.max_retries:
                time.sleep(30)  # 固定等待时间

        error_msg = f"网格 {grid_idx + 1} 在 {self.max_retries} 次尝试后仍然失败"
        self.logger.error(error_msg)
        return False, error_msg