313 lines
11 KiB
Python
313 lines
11 KiB
Python
import os
|
||
from collections import defaultdict
|
||
|
||
import cv2
|
||
import numpy as np
|
||
from sahi import AutoDetectionModel
|
||
from sahi.predict import get_sliced_prediction
|
||
from ultralytics import YOLO
|
||
|
||
from config import DEVICE, DEFAULT_IOU, DEFAULT_MIN_SIZE, DEFAULT_POS_THRESH, SLICE_RULES, DEFAULT_CONF
|
||
|
||
|
||
class YOLODetector:
|
||
def __init__(self, model_path, params, type_to_id):
|
||
# 加载YOLO模型
|
||
self.model = YOLO(model_path)
|
||
self.model.to(DEVICE)
|
||
self.class_names = self.model.names
|
||
self.type_to_id = type_to_id
|
||
|
||
self.params = params
|
||
self.enable_primary = params.get("enable_primary", True)
|
||
self.primary_conf = params.get("primary_conf", DEFAULT_CONF) # 初级检测阈值
|
||
self.secondary_conf = params.get("secondary_conf", DEFAULT_CONF) # 次级检测阈值
|
||
self.final_conf = params.get("final_conf", DEFAULT_CONF) # 最终展示阈值
|
||
|
||
# SAHI模型
|
||
self.sahi_model = None
|
||
if params["enable_secondary"]:
|
||
self.sahi_model = AutoDetectionModel.from_pretrained(
|
||
model_type='yolov8',
|
||
model_path=model_path,
|
||
confidence_threshold=self.secondary_conf,
|
||
device=DEVICE
|
||
)
|
||
|
||
# 统计
|
||
self.stats = defaultdict(int)
|
||
|
||
def get_adaptive_slice(self, total_pixels):
|
||
"""自适应切片参数"""
|
||
for pixel_thresh, (size, overlap) in SLICE_RULES:
|
||
if total_pixels > pixel_thresh:
|
||
return size, overlap
|
||
return self.params["slice_size"], self.params["overlap_ratio"]
|
||
|
||
def multi_scale_detect(self, img_path):
|
||
"""多尺度检测（使用模型专属初级阈值）"""
|
||
detections = []
|
||
img = cv2.imread(img_path)
|
||
h, w = img.shape[:2]
|
||
|
||
for scale in self.params["multi_scales"]:
|
||
if scale == 1.0:
|
||
# 原尺度检测
|
||
results = self.model(
|
||
img_path,
|
||
conf=self.primary_conf, # 模型专属初级阈值
|
||
device=DEVICE,
|
||
classes=self.target_ids,
|
||
verbose=False
|
||
)
|
||
else:
|
||
# 缩放检测
|
||
nw, nh = int(w * scale), int(h * scale)
|
||
scaled_img = cv2.resize(img, (nw, nh))
|
||
temp_path = f"temp_scale_{scale}.jpg"
|
||
cv2.imwrite(temp_path, scaled_img)
|
||
|
||
results = self.model(
|
||
temp_path,
|
||
conf=self.primary_conf, # 模型专属初级阈值
|
||
device=DEVICE,
|
||
classes=self.target_ids,
|
||
verbose=False
|
||
)
|
||
os.remove(temp_path)
|
||
|
||
# 解析结果（核心修复：增加对result.boxes为None的判断）
|
||
for result in results:
|
||
# 检查boxes是否存在且非空
|
||
if result.boxes is None:
|
||
continue
|
||
for box in result.boxes:
|
||
bbox = box.xyxy[0].tolist()
|
||
if scale != 1.0:
|
||
bbox = [coord / scale for coord in bbox]
|
||
|
||
detections.append({
|
||
"box": bbox,
|
||
"conf": box.conf[0].item(),
|
||
"class": box.cls[0].item(),
|
||
"class_name": self.class_names[int(box.cls[0])],
|
||
"source": "primary"
|
||
})
|
||
|
||
return detections
|
||
|
||
def primary_detect(self, img_path):
|
||
"""初次检测（使用模型专属初级阈值）- 新增enable_primary判断"""
|
||
# 新增：如果禁用一级检测，直接返回空列表
|
||
if not self.enable_primary:
|
||
self.stats["primary"] = 0
|
||
print(" 一级检测已禁用，跳过初级检测")
|
||
return []
|
||
|
||
if self.params["enable_multi_scale"]:
|
||
detections = self.multi_scale_detect(img_path)
|
||
else:
|
||
results = self.model(
|
||
img_path,
|
||
conf=self.primary_conf, # 模型专属初级阈值
|
||
device=DEVICE,
|
||
classes=self.target_ids,
|
||
verbose=False
|
||
)
|
||
# 解析结果（核心修复：增加对result.boxes为None的判断）
|
||
detections = []
|
||
for result in results:
|
||
# 检查boxes是否存在且非空
|
||
if result.boxes is None:
|
||
continue
|
||
for box in result.boxes:
|
||
detections.append({
|
||
"box": box.xyxy[0].tolist(),
|
||
"conf": box.conf[0].item(),
|
||
"class": box.cls[0].item(),
|
||
"class_name": self.class_names[int(box.cls[0])],
|
||
"source": "primary"
|
||
})
|
||
|
||
self.stats["primary"] = len(detections)
|
||
return detections
|
||
|
||
def secondary_detect(self, img_path):
|
||
"""SAHI切片检测（已在初始化时使用模型专属次级阈值）"""
|
||
if not self.params["enable_secondary"] or not self.sahi_model:
|
||
return []
|
||
|
||
img = cv2.imread(img_path)
|
||
h, w = img.shape[:2]
|
||
total_pixels = w * h
|
||
slice_size, overlap = self.get_adaptive_slice(total_pixels)
|
||
|
||
# SAHI切片预测
|
||
sliced_results = get_sliced_prediction(
|
||
img_path,
|
||
self.sahi_model,
|
||
slice_height=slice_size,
|
||
slice_width=slice_size,
|
||
overlap_height_ratio=overlap,
|
||
overlap_width_ratio=overlap,
|
||
verbose=0
|
||
)
|
||
|
||
detections = []
|
||
for obj in sliced_results.object_prediction_list:
|
||
if self.target_ids and obj.category.id not in self.target_ids:
|
||
continue
|
||
|
||
bbox = obj.bbox.to_xyxy()
|
||
bw, bh = bbox[2] - bbox[0], bbox[3] - bbox[1]
|
||
|
||
if bw >= DEFAULT_MIN_SIZE and bh >= DEFAULT_MIN_SIZE:
|
||
detections.append({
|
||
"box": bbox,
|
||
"conf": obj.score.value,
|
||
"class": obj.category.id,
|
||
"class_name": obj.category.name,
|
||
"source": "secondary"
|
||
})
|
||
|
||
self.stats["secondary"] = len(detections)
|
||
return detections
|
||
|
||
@staticmethod
|
||
def calculate_iou(box1, box2):
|
||
"""计算IoU"""
|
||
x11, y11, x21, y21 = box1
|
||
x12, y12, x22, y22 = box2
|
||
|
||
inter_x1 = max(x11, x12)
|
||
inter_y1 = max(y11, y12)
|
||
inter_x2 = min(x21, x22)
|
||
inter_y2 = min(y21, y22)
|
||
|
||
inter_area = max(0, inter_x2 - inter_x1) * max(0, inter_y2 - inter_y1)
|
||
area1 = (x21 - x11) * (y21 - y11)
|
||
area2 = (x22 - x12) * (y22 - y12)
|
||
union_area = area1 + area2 - inter_area
|
||
|
||
return inter_area / union_area if union_area > 0 else 0
|
||
|
||
def merge_detections(self, primary_dets, secondary_dets):
|
||
"""融合检测结果"""
|
||
if not primary_dets:
|
||
return secondary_dets
|
||
if not secondary_dets:
|
||
return primary_dets
|
||
|
||
# 加权置信度
|
||
all_dets = []
|
||
for det in primary_dets:
|
||
det["weighted_conf"] = det["conf"] * self.params["weight_primary"]
|
||
all_dets.append(det)
|
||
for det in secondary_dets:
|
||
det["weighted_conf"] = det["conf"] * self.params["weight_secondary"]
|
||
all_dets.append(det)
|
||
|
||
# 按类别分组融合
|
||
class_groups = defaultdict(list)
|
||
for det in all_dets:
|
||
class_groups[det["class"]].append(det)
|
||
|
||
merged = []
|
||
for cls_id, cls_dets in class_groups.items():
|
||
cls_dets.sort(key=lambda x: x["weighted_conf"], reverse=True)
|
||
suppressed = [False] * len(cls_dets)
|
||
|
||
for i in range(len(cls_dets)):
|
||
if suppressed[i]:
|
||
continue
|
||
merged.append(cls_dets[i])
|
||
for j in range(i + 1, len(cls_dets)):
|
||
if not suppressed[j] and self.calculate_iou(cls_dets[i]["box"], cls_dets[j]["box"]) > DEFAULT_IOU:
|
||
suppressed[j] = True
|
||
|
||
self.stats["merged"] = len(merged)
|
||
return merged
|
||
|
||
def post_process(self, detections):
|
||
"""后处理（使用模型专属最终阈值）"""
|
||
# 置信度过滤：模型专属最终阈值
|
||
filtered = [det for det in detections if det["conf"] >= self.final_conf]
|
||
|
||
# 位置去重
|
||
final_dets = []
|
||
for curr_det in filtered:
|
||
curr_cx = (curr_det["box"][0] + curr_det["box"][2]) / 2
|
||
curr_cy = (curr_det["box"][1] + curr_det["box"][3]) / 2
|
||
curr_cls = curr_det["class"]
|
||
duplicate = False
|
||
|
||
for idx, exist_det in enumerate(final_dets):
|
||
if exist_det["class"] != curr_cls:
|
||
continue
|
||
|
||
exist_cx = (exist_det["box"][0] + exist_det["box"][2]) / 2
|
||
exist_cy = (exist_det["box"][1] + exist_det["box"][3]) / 2
|
||
dist = np.sqrt((curr_cx - exist_cx) **2 + (curr_cy - exist_cy)** 2)
|
||
|
||
if dist < DEFAULT_POS_THRESH:
|
||
duplicate = True
|
||
if curr_det["conf"] > exist_det["conf"]:
|
||
final_dets[idx] = curr_det
|
||
break
|
||
|
||
if not duplicate:
|
||
final_dets.append(curr_det)
|
||
|
||
self.stats["final"] = len(final_dets)
|
||
return final_dets
|
||
|
||
def format_results(self, detections):
|
||
"""格式化结果"""
|
||
formatted = []
|
||
for det in detections:
|
||
x1, y1, x2, y2 = det["box"]
|
||
formatted.append({
|
||
"type": det["class_name"],
|
||
"size": [int(round(x2 - x1)), int(round(y2 - y1))],
|
||
"leftTopPoint": [int(round(x1)), int(round(y1))],
|
||
"score": round(det["conf"], 4),
|
||
})
|
||
return formatted
|
||
|
||
def get_detection_stats(self):
|
||
"""获取检测统计信息"""
|
||
return dict(self.stats)
|
||
|
||
def detect(self, img_path, target_types=None):
|
||
"""完整检测流程"""
|
||
# 重置统计
|
||
self.stats = defaultdict(int)
|
||
|
||
# 设置目标类别
|
||
if target_types:
|
||
self.target_ids = [self.type_to_id[cls] for cls in target_types if cls in self.type_to_id]
|
||
else:
|
||
self.target_ids = None
|
||
|
||
# 执行检测
|
||
primary_dets = self.primary_detect(img_path)
|
||
print(f" 初级检测后: {self.stats['primary']} 个目标")
|
||
|
||
if self.params["enable_secondary"]:
|
||
secondary_dets = self.secondary_detect(img_path)
|
||
print(f" 次级检测后: {self.stats['secondary']} 个目标")
|
||
merged_dets = self.merge_detections(primary_dets, secondary_dets)
|
||
print(f" 融合去重后: {self.stats['merged']} 个目标")
|
||
else:
|
||
merged_dets = primary_dets
|
||
print(f" 次级检测未启用")
|
||
|
||
# 后处理
|
||
processed_dets = self.post_process(merged_dets)
|
||
print(f" 过滤低置信度后: {self.stats['final']} 个目标")
|
||
|
||
print(" 最终检测目标详情：")
|
||
for idx, det in enumerate(processed_dets, 1):
|
||
print(f" 目标{idx} - 类型：{det['class_name']}，置信度：{det['conf']:.4f}")
|
||
|
||
return self.format_results(processed_dets) |