|
@@ -3,7 +3,6 @@ import onnxruntime as ort
|
|
|
import cv2
|
|
|
import os
|
|
|
import argparse
|
|
|
-import pickle
|
|
|
from glob import glob
|
|
|
from typing import List, Dict, Optional
|
|
|
import time
|
|
@@ -46,8 +45,7 @@ class ONNXDetector:
|
|
|
self.model_path = r"D:\PythonProject\R360-UAVmodelTool\src\models\UAV-250411.onnx"
|
|
|
else:
|
|
|
self.model_path = r"D:\PythonProject\R360-UAVmodelTool\src\models\uav_and_bird.onnx"
|
|
|
- print(f"模型类型: {self.model_type}")
|
|
|
- print(f"模型路径: {self.model_path}")
|
|
|
+ # 模型初始化完成
|
|
|
|
|
|
# 初始化ONNX会话
|
|
|
so = ort.SessionOptions()
|
|
@@ -74,15 +72,12 @@ class ONNXDetector:
|
|
|
self.output_name = 'conv2d_308.tmp_1'
|
|
|
|
|
|
# 类别映射
|
|
|
- self.class_names = {0: 'UAV'}
|
|
|
+ if self.model_type == 'uav_and_bird':
|
|
|
+ self.class_names = {0: 'Bird', 1: 'UAV'}
|
|
|
+ else:
|
|
|
+ self.class_names = {0: 'UAV'}
|
|
|
|
|
|
- # 加载误报特征库
|
|
|
- self.false_positive_features = {}
|
|
|
- try:
|
|
|
- with open(os.path.join(os.path.dirname(__file__), 'false_positive_features.pkl'), 'rb') as f:
|
|
|
- self.false_positive_features = pickle.load(f)
|
|
|
- except FileNotFoundError:
|
|
|
- print("未找到误报特征库文件,跳过加载")
|
|
|
+ # 误报特征库相关代码已移除
|
|
|
|
|
|
# 创建输出目录
|
|
|
output_base = os.path.join(os.path.dirname(__file__), '..', '..', 'Output')
|
|
@@ -109,7 +104,7 @@ class ONNXDetector:
|
|
|
if self.use_cuda:
|
|
|
self.gpu_frame.upload(cv2.imread(image_path))
|
|
|
if self.gpu_frame.empty():
|
|
|
- print(f"无法加载图像:{image_path}")
|
|
|
+ # 图像加载失败
|
|
|
return None
|
|
|
|
|
|
# GPU预处理流水线
|
|
@@ -123,7 +118,7 @@ class ONNXDetector:
|
|
|
# CPU回退路径
|
|
|
image_orig = cv2.imread(image_path)
|
|
|
if image_orig is None:
|
|
|
- print(f"无法加载图像:{image_path}")
|
|
|
+ # 图像加载失败
|
|
|
return None
|
|
|
|
|
|
self.orig_h, self.orig_w = image_orig.shape[:2]
|
|
@@ -165,83 +160,60 @@ class ONNXDetector:
|
|
|
def inference(self, input_data: np.ndarray):
|
|
|
if self.model_type == 'uav_and_bird':
|
|
|
# 返回所有输出
|
|
|
- return self.session.run(None, {self.input_name: input_data})
|
|
|
+ outputs = self.session.run(None, {self.input_name: input_data})
|
|
|
+ # 模型推理完成
|
|
|
+ return outputs
|
|
|
elif self.model_type in ['Anti_UAV', 'UAV-250411']:
|
|
|
scale_factor = np.array([[1, 1]], dtype=np.float32)
|
|
|
- return self.session.run(
|
|
|
+ output = self.session.run(
|
|
|
[self.output_name],
|
|
|
{self.input_name: input_data, 'scale_factor': scale_factor}
|
|
|
)[0]
|
|
|
+ # 模型推理完成
|
|
|
+ return output
|
|
|
else:
|
|
|
- return self.session.run(
|
|
|
+ output = self.session.run(
|
|
|
[self.output_name],
|
|
|
{self.input_name: input_data}
|
|
|
)[0]
|
|
|
+ # 模型推理完成
|
|
|
+ return output
|
|
|
|
|
|
- def _is_false_positive(self, roi: np.ndarray) -> bool:
|
|
|
- """使用ORB特征匹配验证是否误报"""
|
|
|
- if not self.false_positive_features:
|
|
|
- return False
|
|
|
-
|
|
|
- # 初始化特征检测器
|
|
|
- detector = cv2.ORB_create()
|
|
|
- _, des = detector.detectAndCompute(roi, None)
|
|
|
-
|
|
|
- # 与特征库进行匹配
|
|
|
- for fp_feature in self.false_positive_features.values():
|
|
|
- if des is None or fp_feature['features'] is None:
|
|
|
- continue
|
|
|
-
|
|
|
- # 使用FLANN匹配器
|
|
|
- flann = cv2.FlannBasedMatcher(dict(algorithm=6, table_number=6), dict())
|
|
|
- matches = flann.knnMatch(des, fp_feature['features'], k=2)
|
|
|
-
|
|
|
- # 处理空匹配情况
|
|
|
- if not matches:
|
|
|
- continue
|
|
|
-
|
|
|
- # 安全验证匹配结果
|
|
|
- good_matches = []
|
|
|
- for match_group in matches:
|
|
|
- # 确保match_group有足够元素防止索引错误
|
|
|
- if len(match_group) < 2:
|
|
|
- continue
|
|
|
- m, n = match_group
|
|
|
- if m.distance < 0.7 * n.distance:
|
|
|
- good_matches.append(m)
|
|
|
-
|
|
|
- # 计算优质匹配数量
|
|
|
- if len(good_matches) > 15: # 匹配阈值
|
|
|
- return True
|
|
|
- return False
|
|
|
+ # ORB特征匹配误报检测方法已移除
|
|
|
|
|
|
def postprocess(self, detections, image_orig: np.ndarray, image_path: str) -> tuple:
|
|
|
+ # 开始后处理
|
|
|
+
|
|
|
valid_detections = 0
|
|
|
img_out = image_orig.copy()
|
|
|
detections_list = []
|
|
|
+
|
|
|
if self.model_type == 'Anti_UAV':
|
|
|
# Anti_UAV模型的后处理逻辑
|
|
|
keep_idx = (detections[:, 1] > self.confThreshold)
|
|
|
detections = detections[keep_idx]
|
|
|
+
|
|
|
if len(detections) == 0:
|
|
|
return 0, img_out, []
|
|
|
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])
|
|
|
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)
|
|
|
bbox_area = (x2 - x1) * (y2 - y1)
|
|
|
image_area = self.orig_w * self.orig_h
|
|
|
+
|
|
|
if bbox_area / image_area > self.max_bbox_ratio:
|
|
|
continue
|
|
|
roi = image_orig[y1:y2, x1:x2]
|
|
|
if roi.size == 0:
|
|
|
continue
|
|
|
- if self._is_false_positive(roi):
|
|
|
- continue
|
|
|
+ # 误报过滤已移除
|
|
|
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)
|
|
|
label = f'{self.class_names[class_id]} {confidence:.2f}'
|
|
@@ -260,25 +232,28 @@ class ONNXDetector:
|
|
|
# UAV-250411模型的后处理逻辑
|
|
|
keep_idx = (detections[:, 1] > self.confThreshold)
|
|
|
detections = detections[keep_idx]
|
|
|
+
|
|
|
if len(detections) == 0:
|
|
|
return 0, img_out, []
|
|
|
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])
|
|
|
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)
|
|
|
bbox_area = (x2 - x1) * (y2 - y1)
|
|
|
image_area = self.orig_w * self.orig_h
|
|
|
+
|
|
|
if bbox_area / image_area > self.max_bbox_ratio:
|
|
|
continue
|
|
|
roi = image_orig[y1:y2, x1:x2]
|
|
|
if roi.size == 0:
|
|
|
continue
|
|
|
- if self._is_false_positive(roi):
|
|
|
- continue
|
|
|
+ # 误报过滤已移除
|
|
|
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)
|
|
|
label = f'{self.class_names[class_id]} {confidence:.2f}'
|
|
@@ -294,58 +269,136 @@ class ONNXDetector:
|
|
|
'orig_h': self.orig_h
|
|
|
})
|
|
|
elif self.model_type == 'uav_and_bird':
|
|
|
- # Netron顺序: [bbox20, cls20, obj20, bbox40, cls40, obj40, bbox80, cls80, obj80]
|
|
|
- all_boxes, all_scores, all_classes = [], [], []
|
|
|
+ # 根据用户提供的详细说明,模型输出9个张量,按3个尺度分组:
|
|
|
+ # 小尺度 (20x20): 输出0(bbox), 输出1(cls), 输出2(quality)
|
|
|
+ # 中等尺度 (40x40): 输出3(bbox), 输出4(cls), 输出5(quality)
|
|
|
+ # 大尺度 (80x80): 输出6(bbox), 输出7(cls), 输出8(quality)
|
|
|
+ if len(detections) != 9:
|
|
|
+ return 0, img_out, []
|
|
|
+
|
|
|
+ # 定义尺度信息:(bbox_idx, cls_idx, quality_idx, grid_size, stride)
|
|
|
scales = [
|
|
|
- (0, 1, 2, 20, 20),
|
|
|
- (3, 4, 5, 40, 40),
|
|
|
- (6, 7, 8, 80, 80),
|
|
|
+ (0, 1, 2, 20, 32), # 小尺度特征图,用于检测大物体
|
|
|
+ (3, 4, 5, 40, 16), # 中等尺度特征图,用于检测中等物体
|
|
|
+ (6, 7, 8, 80, 8), # 大尺度特征图,用于检测小物体
|
|
|
]
|
|
|
- for bbox_idx, cls_idx, obj_idx, H, W in scales:
|
|
|
- bbox_pred = detections[bbox_idx][0] # (68, H, W)
|
|
|
- cls_pred = detections[cls_idx][0] # (2, H, W)
|
|
|
- obj_pred = detections[obj_idx][0] # (1, H, W)
|
|
|
- bbox_pred = bbox_pred.reshape(68, -1).T # (H*W, 68)
|
|
|
- cls_pred = cls_pred.reshape(2, -1).T # (H*W, 2)
|
|
|
- obj_pred = obj_pred.reshape(-1) # (H*W,)
|
|
|
- scores = obj_pred[:, None] * (1 / (1 + np.exp(-cls_pred)))
|
|
|
- max_scores = np.max(scores, axis=1)
|
|
|
- class_ids = np.argmax(scores, axis=1)
|
|
|
+
|
|
|
+ all_boxes, all_scores, all_classes = [], [], []
|
|
|
+
|
|
|
+ for scale_idx, (bbox_idx, cls_idx, quality_idx, grid_size, stride) in enumerate(scales):
|
|
|
+ # 获取当前尺度的输出
|
|
|
+ bbox_output = detections[bbox_idx] # (1, 68, H, W)
|
|
|
+ cls_output = detections[cls_idx] # (1, 2, H, W)
|
|
|
+ quality_output = detections[quality_idx] # (1, 1, H, W)
|
|
|
+
|
|
|
+ # 移除batch维度
|
|
|
+ bbox_pred = bbox_output[0] # (68, H, W)
|
|
|
+ cls_pred = cls_output[0] # (2, H, W)
|
|
|
+ quality_pred = quality_output[0] # (1, H, W)
|
|
|
+
|
|
|
+ H, W = bbox_pred.shape[1], bbox_pred.shape[2]
|
|
|
+
|
|
|
+ # 创建网格坐标 - 添加0.5偏移表示像素中心
|
|
|
+ grid_y, grid_x = np.meshgrid(np.arange(H), np.arange(W), indexing='ij')
|
|
|
+ grid_x = (grid_x + 0.5).flatten() # (H*W,) - 像素中心偏移
|
|
|
+ grid_y = (grid_y + 0.5).flatten() # (H*W,) - 像素中心偏移
|
|
|
+
|
|
|
+ # 重塑张量为 (H*W, channels)
|
|
|
+ bbox_pred = bbox_pred.reshape(68, -1).T # (H*W, 68)
|
|
|
+ cls_pred = cls_pred.reshape(2, -1).T # (H*W, 2)
|
|
|
+ quality_pred = quality_pred.reshape(-1) # (H*W,)
|
|
|
+
|
|
|
+ # 解码边界框 - Distribution Focal Loss解码
|
|
|
+ # 68个通道 = 4个边界 * 17个分布值
|
|
|
+ bbox_pred = bbox_pred.reshape(-1, 4, 17) # (H*W, 4, 17)
|
|
|
+
|
|
|
+ # 对每个边界的17个值应用softmax
|
|
|
+ bbox_pred_softmax = np.exp(bbox_pred) / np.sum(np.exp(bbox_pred), axis=2, keepdims=True)
|
|
|
+
|
|
|
+ # 创建范围向量 [0, 1, 2, ..., 16]
|
|
|
+ range_vector = np.arange(17).reshape(1, 1, 17)
|
|
|
+
|
|
|
+ # 计算期望距离值
|
|
|
+ distances = np.sum(bbox_pred_softmax * range_vector, axis=2) # (H*W, 4)
|
|
|
+ d_left, d_top, d_right, d_bottom = distances[:, 0], distances[:, 1], distances[:, 2], distances[:, 3]
|
|
|
+
|
|
|
+ # 计算最终边界框坐标
|
|
|
+ x1 = (grid_x - d_left) * stride
|
|
|
+ y1 = (grid_y - d_top) * stride
|
|
|
+ x2 = (grid_x + d_right) * stride
|
|
|
+ y2 = (grid_y + d_bottom) * stride
|
|
|
+
|
|
|
+ # 组合边界框
|
|
|
+ boxes = np.stack([x1, y1, x2, y2], axis=1) # (H*W, 4)
|
|
|
+
|
|
|
+ # 计算最终置信度分数
|
|
|
+ # cls_pred已经是sigmoid激活后的结果,quality_pred是clip后的结果
|
|
|
+ final_scores = cls_pred * quality_pred[:, None] # (H*W, 2)
|
|
|
+
|
|
|
+ # 获取最大置信度和对应类别
|
|
|
+ max_scores = np.max(final_scores, axis=1) # (H*W,)
|
|
|
+ class_ids = np.argmax(final_scores, axis=1) # (H*W,)
|
|
|
+
|
|
|
+ # 置信度过滤
|
|
|
mask = max_scores > self.confThreshold
|
|
|
- boxes = bbox_pred[mask]
|
|
|
- scores = max_scores[mask]
|
|
|
- classes = class_ids[mask]
|
|
|
- all_boxes.append(boxes)
|
|
|
- all_scores.append(scores)
|
|
|
- all_classes.append(classes)
|
|
|
+
|
|
|
+ if np.sum(mask) > 0:
|
|
|
+ filtered_boxes = boxes[mask]
|
|
|
+ filtered_scores = max_scores[mask]
|
|
|
+ filtered_classes = class_ids[mask]
|
|
|
+
|
|
|
+ all_boxes.append(filtered_boxes)
|
|
|
+ all_scores.append(filtered_scores)
|
|
|
+ all_classes.append(filtered_classes)
|
|
|
+
|
|
|
+ # 检查是否有有效检测
|
|
|
if len(all_boxes) == 0 or all([len(b) == 0 for b in all_boxes]):
|
|
|
return 0, img_out, []
|
|
|
- boxes = np.concatenate(all_boxes, axis=0)
|
|
|
- scores = np.concatenate(all_scores, axis=0)
|
|
|
- classes = np.concatenate(all_classes, axis=0)
|
|
|
- # 坐标缩放
|
|
|
+
|
|
|
+ # 合并所有尺度的检测结果
|
|
|
+ final_boxes = np.concatenate(all_boxes, axis=0)
|
|
|
+ final_scores = np.concatenate(all_scores, axis=0)
|
|
|
+ final_classes = np.concatenate(all_classes, axis=0)
|
|
|
+
|
|
|
+ # 坐标缩放到原图尺寸
|
|
|
ratioh = self.orig_h / self.input_size[1]
|
|
|
ratiow = self.orig_w / self.input_size[0]
|
|
|
- boxes[:, [0, 2]] *= ratiow
|
|
|
- boxes[:, [1, 3]] *= ratioh
|
|
|
- # NMS
|
|
|
- keep = self.nms(boxes, scores[:, None], self.confThreshold, 0.4)
|
|
|
+
|
|
|
+ final_boxes[:, [0, 2]] *= ratiow # x坐标
|
|
|
+ final_boxes[:, [1, 3]] *= ratioh # y坐标
|
|
|
+
|
|
|
+ # NMS处理
|
|
|
+ keep = self.nms(final_boxes, final_scores[:, None], self.confThreshold, 0.4)
|
|
|
+
|
|
|
+ # 处理最终检测结果
|
|
|
for idx in keep:
|
|
|
- x1, y1, x2, y2 = boxes[idx][:4].astype(int)
|
|
|
- confidence = scores[idx]
|
|
|
- class_id = int(classes[idx])
|
|
|
+ x1, y1, x2, y2 = final_boxes[idx].astype(int)
|
|
|
+ confidence = final_scores[idx]
|
|
|
+ class_id = int(final_classes[idx])
|
|
|
+
|
|
|
+ # 确保坐标在图像范围内
|
|
|
+ x1 = max(0, min(x1, self.orig_w - 1))
|
|
|
+ y1 = max(0, min(y1, self.orig_h - 1))
|
|
|
+ x2 = max(0, min(x2, self.orig_w - 1))
|
|
|
+ y2 = max(0, min(y2, self.orig_h - 1))
|
|
|
+
|
|
|
bbox_area = (x2 - x1) * (y2 - y1)
|
|
|
image_area = self.orig_w * self.orig_h
|
|
|
+
|
|
|
if bbox_area / image_area > self.max_bbox_ratio:
|
|
|
continue
|
|
|
+
|
|
|
+ if x2 <= x1 or y2 <= y1:
|
|
|
+ continue
|
|
|
+
|
|
|
roi = image_orig[y1:y2, x1:x2]
|
|
|
if roi.size == 0:
|
|
|
continue
|
|
|
- if self._is_false_positive(roi):
|
|
|
- continue
|
|
|
+
|
|
|
+ # 误报过滤已移除
|
|
|
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)
|
|
|
- label = f'class {class_id} {confidence:.2f}'
|
|
|
+ label = f'{self.class_names[class_id]} {confidence:.2f}'
|
|
|
cv2.rectangle(img_out, (x1, y1), (x2, y2), (255, 0, 0), 4)
|
|
|
cv2.putText(img_out, label, (x1, y1 - 10),
|
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
|
|
@@ -357,6 +410,7 @@ class ONNXDetector:
|
|
|
'orig_w': self.orig_w,
|
|
|
'orig_h': self.orig_h
|
|
|
})
|
|
|
+
|
|
|
return valid_detections, img_out, detections_list
|
|
|
|
|
|
def process_image(self, image_path: str) -> int:
|
|
@@ -384,8 +438,11 @@ class ONNXDetector:
|
|
|
# 根据设置保存图片
|
|
|
if valid_count > 0 or self.save_empty:
|
|
|
cv2.imwrite(output_path, processed_img)
|
|
|
+ # 输出图像已保存
|
|
|
else:
|
|
|
output_path = None
|
|
|
+ # 未保存图像 (无检测且save_empty=False)
|
|
|
+
|
|
|
self.image_count += 1
|
|
|
|
|
|
# 记录检测信息
|
|
@@ -393,6 +450,7 @@ class ONNXDetector:
|
|
|
record['detections'] = detections_list
|
|
|
|
|
|
self.detection_records.append(record)
|
|
|
+ # 处理完成
|
|
|
return valid_count
|
|
|
|
|
|
def get_model_type(model_path: str) -> str:
|
|
@@ -422,7 +480,8 @@ def get_model_type(model_path: str) -> str:
|
|
|
return 'uav_and_bird'
|
|
|
return 'unknown'
|
|
|
except Exception as e:
|
|
|
- print(f"模型类型识别失败: {e}")
|
|
|
+ # 模型类型识别失败
|
|
|
+ pass
|
|
|
return 'unknown'
|
|
|
|
|
|
# 命令行接口
|
|
@@ -432,13 +491,13 @@ if __name__ == '__main__':
|
|
|
parser.add_argument('--threshold', type=float, default=0.5, help='检测置信度阈值')
|
|
|
parser.add_argument('--output', type=str, default=None, help='输出目录路径,默认为输入目录名+_results')
|
|
|
parser.add_argument('--max-bbox-ratio', type=float, default=0.05,
|
|
|
- help='检测框最大面积比例阈值,默认0.05')
|
|
|
+ help='检测框最大面积比例阈值,默认0.05')
|
|
|
parser.add_argument('--save-empty', action='store_true',
|
|
|
- help='是否保存未检测到目标的图片')
|
|
|
+ help='是否保存未检测到目标的图片')
|
|
|
parser.add_argument('--gui', action='store_true',
|
|
|
- help='启用图形界面选择输入目录')
|
|
|
+ help='启用图形界面选择输入目录')
|
|
|
parser.add_argument('--model-type', type=str, choices=['Anti_UAV', 'UAV-250411', 'uav_and_bird'],
|
|
|
- help='指定模型类型,不指定则自动识别')
|
|
|
+ help='指定模型类型,不指定则自动识别')
|
|
|
parser.add_argument('--model-path', type=str, help='指定模型路径,不指定则根据模型类型自动选择')
|
|
|
args = parser.parse_args()
|
|
|
|
|
@@ -454,7 +513,7 @@ if __name__ == '__main__':
|
|
|
)
|
|
|
input_dir = detector.select_input_directory()
|
|
|
if not input_dir:
|
|
|
- print("未选择目录,程序退出")
|
|
|
+ # 未选择目录,程序退出
|
|
|
exit()
|
|
|
args.input = input_dir
|
|
|
|
|
@@ -471,7 +530,7 @@ if __name__ == '__main__':
|
|
|
|
|
|
def process_single(image_path: str):
|
|
|
detections = detector.process_image(image_path)
|
|
|
- print(f'处理 {os.path.basename(image_path)} 完成,检测到 {detections} 个目标')
|
|
|
+ # 图像处理完成
|
|
|
|
|
|
if os.path.isdir(args.input):
|
|
|
total = 0
|
|
@@ -480,11 +539,11 @@ if __name__ == '__main__':
|
|
|
image_files.extend([os.path.join(root, f) for f in files if f.lower().endswith('.jpg')])
|
|
|
for img_file in tqdm(image_files, desc='Processing images'):
|
|
|
total += detector.process_image(img_file)
|
|
|
- print(f'批量处理完成!共检测到 {total} 个目标')
|
|
|
+ # 批量处理完成
|
|
|
|
|
|
# 生成CSV报告
|
|
|
csv_path = os.path.join(detector.output_dir, f'detection_report_{detector.model_type}.csv')
|
|
|
ReportGenerator(detector).generate_csv(csv_path)
|
|
|
- print(f'CSV报告已生成: {csv_path}')
|
|
|
+ # CSV报告已生成
|
|
|
else:
|
|
|
detections = detector.process_image(args.input)
|