UAV/utils/directory_manager.py

import os
import shutil
import psutil


class DirectoryManager:
    def __init__(self, config):
        """
        初始化目录管理器
        Args:
            config: 配置对象，包含输入和输出目录等信息
        """
        self.config = config

    def clean_output_dir(self):
        """清理输出目录"""
        try:
            if os.path.exists(self.config.output_dir):
                shutil.rmtree(self.config.output_dir)
                print(f"已清理输出目录: {self.config.output_dir}")
            else:
                pass
        except Exception as e:
            print(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'))

            print(f"已创建输出目录结构: {self.config.output_dir}")
        except Exception as e:
            print(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

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

        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)