修改输入目录结构
This commit is contained in:
weixin_46229132
parent
4681f18085
commit
4c460515c5
120
main.py
120
main.py
@ -7,12 +7,13 @@ import psutil
|
||||
import pandas as pd
|
||||
from pathlib import Path
|
||||
|
||||
from filter.cluster_filter import GPSCluster
|
||||
from utils.directory_manager import DirectoryManager
|
||||
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 filter.cluster_filter import GPSCluster
|
||||
from post_pro.conv_obj import ConvertOBJ
|
||||
|
||||
|
||||
@ -47,124 +48,23 @@ class ProcessConfig:
|
||||
class ODM_Plugin:
|
||||
def __init__(self, config: ProcessConfig):
|
||||
self.config = config
|
||||
self.logger = setup_logger(config.output_dir)
|
||||
|
||||
# 检查磁盘空间
|
||||
# TODO 现在输入目录的磁盘空间也需要检查
|
||||
self._check_disk_space()
|
||||
|
||||
# 初始化目录管理器
|
||||
self.dir_manager = DirectoryManager(config)
|
||||
# 清理并重建输出目录
|
||||
if os.path.exists(config.output_dir):
|
||||
self._clean_output_dir()
|
||||
self._setup_output_dirs()
|
||||
|
||||
self.dir_manager.clean_output_dir()
|
||||
self.dir_manager.setup_output_dirs()
|
||||
# 检查磁盘空间
|
||||
self.dir_manager.check_disk_space()
|
||||
# 修改输入目录，符合ODM要求，从这里开始，image_dir就是project_path
|
||||
self._rename_input_dir()
|
||||
self.project_path = self.config.image_dir
|
||||
self.project_path = self.dir_manager.rename_input_dir()
|
||||
|
||||
# 初始化其他组件
|
||||
self.logger = setup_logger(config.output_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数据")
|
||||
|
||||
2
run.py
2
run.py
@ -10,7 +10,7 @@ def parse_args():
|
||||
# parser.add_argument('--image_dir', required=True, help='输入图片目录路径')
|
||||
# parser.add_argument('--output_dir', required=True, help='输出目录路径')
|
||||
parser.add_argument(
|
||||
'--image_dir', default=r'E:\datasets\UAV\134', help='输入图片目录路径')
|
||||
'--image_dir', default=r'E:\datasets\UAV\134\project\images', help='输入图片目录路径')
|
||||
parser.add_argument(
|
||||
'--output_dir', default=r'G:\ODM_output\134', help='输出目录路径')
|
||||
# 可选参数
|
||||
|
||||
118
utils/directory_manager.py
Normal file
118
utils/directory_manager.py
Normal file
@ -0,0 +1,118 @@
|
||||
import os
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
import psutil
|
||||
|
||||
|
||||
class DirectoryManager:
|
||||
def __init__(self, config):
|
||||
"""
|
||||
初始化目录管理器
|
||||
Args:
|
||||
config: 配置对象，包含输入和输出目录等信息
|
||||
"""
|
||||
self.config = config
|
||||
|
||||
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.image_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 * 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)
|
||||
|
||||
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}")
|
||||
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 rename_input_dir(self):
|
||||
"""修改输入目录，符合ODM要求"""
|
||||
image_dir = Path(self.config.image_dir).resolve()
|
||||
# 原目录名和父路径
|
||||
parent_dir = image_dir.parent
|
||||
original_name = image_dir.name
|
||||
|
||||
if not image_dir.exists() or not image_dir.is_dir():
|
||||
raise ValueError(
|
||||
f"Provided path '{image_dir}' is not a valid directory.")
|
||||
|
||||
if image_dir.name == "images" and parent_dir == image_dir.parent:
|
||||
# 如果目录已经是 images，直接返回
|
||||
print(f"输入目录已经是images，无需重命名: {image_dir}")
|
||||
final_images_path = image_dir
|
||||
else:
|
||||
# 新的 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))
|
||||
|
||||
print(f"符合标准输入的文件夹结构已经创建好了，{final_images_path}")
|
||||
|
||||
return final_images_path.parent.parent
|
||||
|
||||
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
|
||||
@ -37,6 +37,7 @@ class DockerRunner:
|
||||
"--max-concurrency", "15",
|
||||
"--force-gps",
|
||||
"--split-overlap", "0",
|
||||
"--rerun-all"
|
||||
]
|
||||
|
||||
# 运行容器
|
||||
|
||||
Loading…
Reference in New Issue
Block a user