修改输入目录结构

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.dir_manager.clean_output_dir()
-            self._clean_output_dir()
+        self.dir_manager.setup_output_dirs()
-        self._setup_output_dirs()
+        # 检查磁盘空间
-
+        self.dir_manager.check_disk_space()
         # 修改输入目录，符合ODM要求，从这里开始，image_dir就是project_path
-        self._rename_input_dir()
+        self.project_path = self.dir_manager.rename_input_dir()
-        self.project_path = self.config.image_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数据")

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='输出目录路径')
     # 可选参数

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

utils/docker_runner.py

@@ -37,6 +37,7 @@ class DockerRunner:
                 "--max-concurrency", "15",
                 "--force-gps",
                 "--split-overlap", "0",
+                "--rerun-all"
             ]
 
             # 运行容器