Selaa lähdekoodia

refactor(core/inference): 重构输出目录和文件名生成逻辑

重构了ONNXDetector类的输出目录和文件名生成逻辑,确保输出文件组织更清晰。将输出目录统一放置在项目根目录下的Output文件夹中,并添加时间戳以避免命名冲突。同时,简化了目标文件的命名规则,移除了冗余的时间戳信息。
Hannnk 3 viikkoa sitten
vanhempi
commit
eda38add49
2 muutettua tiedostoa jossa 16 lisäystä ja 9 poistoa
  1. 2 2
      src/analysis/performance_analyzer.py
  2. 14 7
      src/core/inference.py

+ 2 - 2
src/analysis/performance_analyzer.py

@@ -9,8 +9,8 @@ class PerformanceAnalyzer:
     def __init__(self, result_dir: str):
         self.result_dir = result_dir
         self.csv_path = os.path.join(result_dir, 'detection_report.csv')
-        self.targets_dir = os.path.join(result_dir, 'targets_all')
-        self.true_targets_dir = os.path.join(result_dir, 'targets_all_True')
+        self.targets_dir = os.path.join(result_dir, 'targets')
+        self.true_targets_dir = os.path.join(result_dir, 'targets_True')
         
         # 确保目录存在
         os.makedirs(self.true_targets_dir, exist_ok=True)

+ 14 - 7
src/core/inference.py

@@ -21,9 +21,9 @@ class ONNXDetector:
         return folder_path if folder_path else None
 
     def __init__(self, model_path: str = 'D:/PythonProject/Model/Data/models/250411_Anti_UAV.onnx', threshold: float = 0.5, 
-                 output_dir: str = "None", save_empty: bool = False, 
+                 output_dir: str = None, save_empty: bool = False, 
                  max_bbox_ratio: float = 0.5,
-                 input_dir: str = "None"):
+                 input_dir: str = None):
         self.image_count = 0
         self.detection_records = []
         self.model_path = model_path
@@ -63,10 +63,15 @@ class ONNXDetector:
             print("未找到误报特征库文件,跳过加载")
         
         # 创建输出目录
+        output_base = os.path.join(os.path.dirname(__file__), '..', '..', 'Output')
+        os.makedirs(output_base, exist_ok=True)
+        
+        timestamp = datetime.now().strftime('%m%d%H%M%S')
         if output_dir is None and input_dir is not None:
-            self.output_dir = f"{os.path.basename(os.path.normpath(input_dir))}_results"
+            self.output_dir = os.path.join(output_base, f"{os.path.basename(os.path.normpath(input_dir))}_{timestamp}")
         else:
-            self.output_dir = output_dir
+            self.output_dir = os.path.join(output_base, output_dir or f"output_{timestamp}")
+            
         # 创建目标存储目录
         self.targets_dir = os.path.join(self.output_dir, 'targets')
         os.makedirs(self.targets_dir, exist_ok=True)
@@ -221,8 +226,7 @@ class ONNXDetector:
                 continue
                 
             # 生成唯一文件名并保存目标
-            timestamp = datetime.now().strftime('%Y%m%d%H%M%S%f')
-            target_filename = f"{os.path.splitext(os.path.basename(image_path))[0]}_{timestamp}_{valid_detections}.jpg"
+            target_filename = f"{os.path.splitext(os.path.basename(image_path))[0]}_{valid_detections}.jpg"
             cv2.imwrite(os.path.join(self.targets_dir, target_filename), roi)
             
             # 绘制检测框和标签
@@ -258,7 +262,10 @@ class ONNXDetector:
         }
 
         filename = f"{os.path.splitext(os.path.basename(image_path))[0]}.jpg"
-        output_path = os.path.join(self.output_dir, filename)
+        # 创建输出图片目录
+        output_img_dir = os.path.join(self.output_dir, 'imgs')
+        os.makedirs(output_img_dir, exist_ok=True)
+        output_path = os.path.join(output_img_dir, filename)
         
         # 根据设置保存图片
         if valid_count > 0 or self.save_empty: