diff --git a/main.py b/main.py index e7bfb7f..5ba6dad 100644 --- a/main.py +++ b/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数据") diff --git a/run.py b/run.py index 1f3cfef..ec1c4cf 100644 --- a/run.py +++ b/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='输出目录路径') # 可选参数 diff --git a/utils/directory_manager.py b/utils/directory_manager.py new file mode 100644 index 0000000..8683569 --- /dev/null +++ b/utils/directory_manager.py @@ -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 diff --git a/utils/docker_runner.py b/utils/docker_runner.py index 19c2d30..d510410 100644 --- a/utils/docker_runner.py +++ b/utils/docker_runner.py @@ -37,6 +37,7 @@ class DockerRunner: "--max-concurrency", "15", "--force-gps", "--split-overlap", "0", + "--rerun-all" ] # 运行容器