|
@@ -20,30 +20,34 @@ class ONNXDetector:
|
|
folder_path = filedialog.askdirectory(title='选择输入目录')
|
|
folder_path = filedialog.askdirectory(title='选择输入目录')
|
|
return folder_path if folder_path else None
|
|
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,
|
|
|
|
|
|
+ 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, # 已有该参数
|
|
|
|
|
|
+ max_bbox_ratio: float = 0.5,
|
|
input_dir: str = "None"):
|
|
input_dir: str = "None"):
|
|
self.image_count = 0
|
|
self.image_count = 0
|
|
self.detection_records = []
|
|
self.detection_records = []
|
|
self.model_path = model_path
|
|
self.model_path = model_path
|
|
self.input_dir = input_dir
|
|
self.input_dir = input_dir
|
|
|
|
+ self.save_empty = save_empty
|
|
|
|
+ self.threshold = threshold
|
|
|
|
+ self.confThreshold = threshold
|
|
|
|
|
|
# 初始化ONNX会话
|
|
# 初始化ONNX会话
|
|
|
|
+ so = ort.SessionOptions()
|
|
|
|
+ so.log_severity_level = 3
|
|
self.providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
|
|
self.providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
|
|
- self.session = ort.InferenceSession(model_path, providers=self.providers)
|
|
|
|
|
|
+ self.session = ort.InferenceSession(model_path, so, providers=self.providers)
|
|
|
|
|
|
- # 获取模型输入尺寸
|
|
|
|
- input_dims = self.session.get_inputs()[0].shape
|
|
|
|
- self.input_shape = tuple([dim if isinstance(dim, int) else 1 for dim in input_dims])
|
|
|
|
|
|
+ # 模型参数设置
|
|
|
|
+ self.input_size = (640, 640) # width, height
|
|
|
|
+ self.mean_ = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
|
|
|
+ self.std_ = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
|
|
|
+ self.max_bbox_ratio = max_bbox_ratio
|
|
|
|
|
|
# 从模型路径解析版本号
|
|
# 从模型路径解析版本号
|
|
- self.model_version = model_path.split('-')[-1].split('.')[0] if '-' in model_path else '1.0'
|
|
|
|
- self.input_size = f"{self.input_shape[2]}x{self.input_shape[3]}"
|
|
|
|
|
|
+ self.model_version = os.path.basename(model_path).split('.')[0]
|
|
|
|
|
|
- self.threshold = threshold
|
|
|
|
- self.max_bbox_ratio = max_bbox_ratio # 已有该赋值
|
|
|
|
- self.save_empty = save_empty
|
|
|
|
|
|
+ # 获取输入输出名称
|
|
self.input_name = self.session.get_inputs()[0].name
|
|
self.input_name = self.session.get_inputs()[0].name
|
|
self.output_name = self.session.get_outputs()[0].name
|
|
self.output_name = self.session.get_outputs()[0].name
|
|
|
|
|
|
@@ -82,7 +86,7 @@ class ONNXDetector:
|
|
return None
|
|
return None
|
|
|
|
|
|
# GPU预处理流水线
|
|
# GPU预处理流水线
|
|
- gpu_resized = cv2.cuda.resize(self.gpu_frame, (self.input_shape[3], self.input_shape[2]))
|
|
|
|
|
|
+ gpu_resized = cv2.cuda.resize(self.gpu_frame, self.input_size)
|
|
gpu_rgb = cv2.cuda.cvtColor(gpu_resized, cv2.COLOR_BGR2RGB)
|
|
gpu_rgb = cv2.cuda.cvtColor(gpu_resized, cv2.COLOR_BGR2RGB)
|
|
|
|
|
|
# 下载到CPU进行后续处理
|
|
# 下载到CPU进行后续处理
|
|
@@ -96,14 +100,41 @@ class ONNXDetector:
|
|
return None
|
|
return None
|
|
|
|
|
|
self.orig_h, self.orig_w = image_orig.shape[:2]
|
|
self.orig_h, self.orig_w = image_orig.shape[:2]
|
|
- image_orig = cv2.resize(image_orig, (self.input_shape[3], self.input_shape[2]))
|
|
|
|
|
|
+ image_orig = cv2.resize(image_orig, self.input_size)
|
|
image_orig = cv2.cvtColor(image_orig, cv2.COLOR_BGR2RGB)
|
|
image_orig = cv2.cvtColor(image_orig, cv2.COLOR_BGR2RGB)
|
|
|
|
|
|
# 统一的后处理
|
|
# 统一的后处理
|
|
image = image_orig.astype(np.float32) / 255.0
|
|
image = image_orig.astype(np.float32) / 255.0
|
|
|
|
+ image -= self.mean_[None, None, :]
|
|
|
|
+ image /= self.std_[None, None, :]
|
|
image = np.transpose(image, (2, 0, 1)) # CHW 格式
|
|
image = np.transpose(image, (2, 0, 1)) # CHW 格式
|
|
return np.expand_dims(image, axis=0)
|
|
return np.expand_dims(image, axis=0)
|
|
|
|
|
|
|
|
+ def nms(self, boxes: np.ndarray, scores: np.ndarray, conf_threshold: float, iou_threshold: float) -> List[int]:
|
|
|
|
+ """非极大值抑制"""
|
|
|
|
+ x1 = boxes[:, 0]
|
|
|
|
+ y1 = boxes[:, 1]
|
|
|
|
+ x2 = boxes[:, 2]
|
|
|
|
+ y2 = boxes[:, 3]
|
|
|
|
+ areas = (x2 - x1 + 1) * (y2 - y1 + 1)
|
|
|
|
+ order = scores.argsort()[0][::-1]
|
|
|
|
+ keep = []
|
|
|
|
+
|
|
|
|
+ while order.size > 0:
|
|
|
|
+ i = order[0]
|
|
|
|
+ keep.append(i)
|
|
|
|
+ xx1 = np.maximum(x1[i], x1[order[1:]])
|
|
|
|
+ yy1 = np.maximum(y1[i], y1[order[1:]])
|
|
|
|
+ xx2 = np.minimum(x2[i], x2[order[1:]])
|
|
|
|
+ yy2 = np.minimum(y2[i], y2[order[1:]])
|
|
|
|
+ w = np.maximum(0.0, xx2 - xx1 + 1)
|
|
|
|
+ h = np.maximum(0.0, yy2 - yy1 + 1)
|
|
|
|
+ inter = w * h
|
|
|
|
+ ovr = inter / (areas[i] + areas[order[1:]] - inter)
|
|
|
|
+ inds = np.where(ovr <= iou_threshold)[0]
|
|
|
|
+ order = order[inds + 1]
|
|
|
|
+ return keep
|
|
|
|
+
|
|
def inference(self, input_data: np.ndarray) -> np.ndarray:
|
|
def inference(self, input_data: np.ndarray) -> np.ndarray:
|
|
"""执行模型推理"""
|
|
"""执行模型推理"""
|
|
scale_factor = np.array([[1, 1]], dtype=np.float32)
|
|
scale_factor = np.array([[1, 1]], dtype=np.float32)
|
|
@@ -149,22 +180,32 @@ class ONNXDetector:
|
|
return True
|
|
return True
|
|
return False
|
|
return False
|
|
|
|
|
|
- def postprocess(self, detections: np.ndarray, image_orig: np.ndarray, image_path: str) -> int:
|
|
|
|
|
|
+ def postprocess(self, detections: np.ndarray, image_orig: np.ndarray, image_path: str) -> tuple:
|
|
"""后处理检测结果"""
|
|
"""后处理检测结果"""
|
|
valid_detections = 0
|
|
valid_detections = 0
|
|
img_out = image_orig.copy()
|
|
img_out = image_orig.copy()
|
|
detections_list = []
|
|
detections_list = []
|
|
|
|
|
|
- for det in detections:
|
|
|
|
- class_id = int(det[0])
|
|
|
|
- confidence = det[1]
|
|
|
|
|
|
+ # 过滤低置信度检测
|
|
|
|
+ keep_idx = (detections[:, 1] > self.confThreshold)
|
|
|
|
+ detections = detections[keep_idx]
|
|
|
|
+
|
|
|
|
+ if len(detections) == 0:
|
|
|
|
+ return 0, img_out, []
|
|
|
|
|
|
- if confidence < self.threshold:
|
|
|
|
- continue
|
|
|
|
|
|
+ # 坐标转换
|
|
|
|
+ ratioh = self.orig_h / self.input_size[1]
|
|
|
|
+ ratiow = self.orig_w / self.input_size[0]
|
|
|
|
+ detections[:, 2:6] *= np.array([ratiow, ratioh, ratiow, ratioh])
|
|
|
|
+
|
|
|
|
+ # NMS处理
|
|
|
|
+ keep = self.nms(detections[:, 2:6], detections[:, 1:2], self.confThreshold, 0.4)
|
|
|
|
+
|
|
|
|
+ for idx in keep:
|
|
|
|
+ class_id = int(detections[idx, 0])
|
|
|
|
+ confidence = detections[idx, 1]
|
|
|
|
+ x1, y1, x2, y2 = detections[idx, 2:6].astype(int)
|
|
|
|
|
|
- # 坐标转换和边界检查
|
|
|
|
- x1, y1, x2, y2 = self._convert_coordinates(det[2:6])
|
|
|
|
-
|
|
|
|
# 计算检测框面积比例
|
|
# 计算检测框面积比例
|
|
bbox_area = (x2 - x1) * (y2 - y1)
|
|
bbox_area = (x2 - x1) * (y2 - y1)
|
|
image_area = self.orig_w * self.orig_h
|
|
image_area = self.orig_w * self.orig_h
|
|
@@ -194,28 +235,12 @@ class ONNXDetector:
|
|
'class_id': class_id,
|
|
'class_id': class_id,
|
|
'confidence': float(confidence),
|
|
'confidence': float(confidence),
|
|
'bbox': [x1, y1, x2, y2],
|
|
'bbox': [x1, y1, x2, y2],
|
|
- 'orig_w': self.orig_w,
|
|
|
|
- 'orig_h': self.orig_h
|
|
|
|
|
|
+ 'orig_w': self.orig_w,
|
|
|
|
+ 'orig_h': self.orig_h
|
|
})
|
|
})
|
|
|
|
|
|
return valid_detections, img_out, detections_list
|
|
return valid_detections, img_out, detections_list
|
|
|
|
|
|
- def _convert_coordinates(self, coords: List[float]) -> tuple:
|
|
|
|
- """将模型输出坐标转换为原始图像尺寸"""
|
|
|
|
- x1 = int(coords[0] * self.orig_w / self.input_shape[3])
|
|
|
|
- y1 = int(coords[1] * self.orig_h / self.input_shape[2])
|
|
|
|
- x2 = int(coords[2] * self.orig_w / self.input_shape[3])
|
|
|
|
- y2 = int(coords[3] * self.orig_h / self.input_shape[2])
|
|
|
|
-
|
|
|
|
- # 边界检查
|
|
|
|
- clamp = lambda val, max_val: max(0, min(val, max_val - 1))
|
|
|
|
- return (
|
|
|
|
- clamp(x1, self.orig_w),
|
|
|
|
- clamp(y1, self.orig_h),
|
|
|
|
- clamp(x2, self.orig_w),
|
|
|
|
- clamp(y2, self.orig_h)
|
|
|
|
- )
|
|
|
|
-
|
|
|
|
def process_image(self, image_path: str) -> int:
|
|
def process_image(self, image_path: str) -> int:
|
|
"""端到端处理单张图像"""
|
|
"""端到端处理单张图像"""
|
|
input_data = self.preprocess(image_path)
|
|
input_data = self.preprocess(image_path)
|
|
@@ -257,8 +282,6 @@ if __name__ == '__main__':
|
|
parser.add_argument('--output', type=str, default=None, help='输出目录路径,默认为输入目录名+_results')
|
|
parser.add_argument('--output', type=str, default=None, help='输出目录路径,默认为输入目录名+_results')
|
|
parser.add_argument('--max-bbox-ratio', type=float, default=0.05,
|
|
parser.add_argument('--max-bbox-ratio', type=float, default=0.05,
|
|
help='检测框最大面积比例阈值,默认0.05')
|
|
help='检测框最大面积比例阈值,默认0.05')
|
|
- parser.add_argument('--report', type=str, choices=['csv', 'excel', 'all'], default='all',
|
|
|
|
- help='输出报告格式: csv|excel|all')
|
|
|
|
parser.add_argument('--save-empty', action='store_true',
|
|
parser.add_argument('--save-empty', action='store_true',
|
|
help='是否保存未检测到目标的图片')
|
|
help='是否保存未检测到目标的图片')
|
|
parser.add_argument('--gui', action='store_true',
|
|
parser.add_argument('--gui', action='store_true',
|
|
@@ -282,7 +305,7 @@ if __name__ == '__main__':
|
|
# 初始化检测器时传递参数
|
|
# 初始化检测器时传递参数
|
|
detector = ONNXDetector(
|
|
detector = ONNXDetector(
|
|
threshold=args.threshold,
|
|
threshold=args.threshold,
|
|
- max_bbox_ratio=args.max_bbox_ratio, # 添加该参数传递
|
|
|
|
|
|
+ max_bbox_ratio=args.max_bbox_ratio,
|
|
output_dir=args.output,
|
|
output_dir=args.output,
|
|
save_empty=args.save_empty,
|
|
save_empty=args.save_empty,
|
|
input_dir=args.input if os.path.isdir(args.input) else None
|
|
input_dir=args.input if os.path.isdir(args.input) else None
|
|
@@ -301,14 +324,9 @@ if __name__ == '__main__':
|
|
total += detector.process_image(img_file)
|
|
total += detector.process_image(img_file)
|
|
print(f'批量处理完成!共检测到 {total} 个目标')
|
|
print(f'批量处理完成!共检测到 {total} 个目标')
|
|
|
|
|
|
- # 生成检测报告
|
|
|
|
- if args.report in ('csv', 'all'):
|
|
|
|
|
|
+ # 生成CSV报告
|
|
csv_path = os.path.join(detector.output_dir, 'detection_report.csv')
|
|
csv_path = os.path.join(detector.output_dir, 'detection_report.csv')
|
|
ReportGenerator(detector).generate_csv(csv_path)
|
|
ReportGenerator(detector).generate_csv(csv_path)
|
|
print(f'CSV报告已生成: {csv_path}')
|
|
print(f'CSV报告已生成: {csv_path}')
|
|
- if args.report in ('excel', 'all'):
|
|
|
|
- excel_path = os.path.join(detector.output_dir, 'detection_report.xlsx')
|
|
|
|
- ReportGenerator(detector).generate_excel(excel_path)
|
|
|
|
- print(f'Excel报告已生成: {excel_path}')
|
|
|
|
else:
|
|
else:
|
|
detections = detector.process_image(args.input)
|
|
detections = detector.process_image(args.input)
|