|
@@ -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:
|