79 lines
2.6 KiB
Python
79 lines
2.6 KiB
Python
![]() |
import os
|
|||
|
import shutil
|
|||
|
import psutil
|
|||
|
|
|||
|
|
|||
|
class DirectoryManager:
|
|||
|
def __init__(self, config):
|
|||
|
"""
|
|||
|
初始化目录管理器
|
|||
|
Args:
|
|||
|
config: 配置对象,包含输入和输出目录等信息
|
|||
|
"""
|
|||
|
self.config = config
|
|||
|
|
|||
|
def clean_output_dir(self):
|
|||
|
"""清理输出目录"""
|
|||
|
try:
|
|||
|
shutil.rmtree(self.config.output_dir)
|
|||
|
print(f"已清理输出目录: {self.config.output_dir}")
|
|||
|
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 * 10
|
|||
|
|
|||
|
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)
|