weijuesen
187 lines
6.4 KiB
Python
187 lines
6.4 KiB
Python
import io
|
|
import os
|
|
from abc import ABC, abstractmethod
|
|
|
|
from PIL import Image
|
|
|
|
|
|
def _label_set(labels) -> set[str]:
|
|
return {label.strip().lower() for label in labels if label.strip()}
|
|
|
|
|
|
def abnormal_probability(detections: list[dict], labels: tuple[str, ...] | list[str]) -> float:
|
|
"""单帧异常概率：取异常类检测的最高置信度，避免多框求和造成虚高。"""
|
|
abnormal = _label_set(labels) - {"healthy"}
|
|
best = 0.0
|
|
for detection in detections:
|
|
class_name = str(detection.get("class_name", "")).strip().lower()
|
|
if class_name in ("", "healthy", "unknown"):
|
|
continue
|
|
if abnormal and class_name not in abnormal:
|
|
continue
|
|
try:
|
|
confidence = float(detection.get("confidence", 0))
|
|
except (TypeError, ValueError):
|
|
confidence = 0.0
|
|
best = max(best, min(1.0, max(0.0, confidence)))
|
|
return best
|
|
|
|
|
|
def detection_status(detections: list[dict], labels: tuple[str, ...] | list[str]) -> str:
|
|
"""空检测或 unknown 类不再被当作 healthy。"""
|
|
if not detections:
|
|
return "unknown"
|
|
abnormal = _label_set(labels) - {"healthy"}
|
|
healthy_seen = False
|
|
unknown_seen = False
|
|
for detection in detections:
|
|
class_name = str(detection.get("class_name", "")).strip().lower()
|
|
if abnormal and class_name in abnormal:
|
|
return "abnormal"
|
|
if not abnormal and class_name not in ("", "healthy", "unknown"):
|
|
return "abnormal"
|
|
if class_name == "healthy":
|
|
healthy_seen = True
|
|
elif class_name in ("", "unknown"):
|
|
unknown_seen = True
|
|
else:
|
|
unknown_seen = True
|
|
if healthy_seen and not unknown_seen:
|
|
return "healthy"
|
|
return "unknown"
|
|
|
|
|
|
class Detector(ABC):
|
|
@abstractmethod
|
|
def detect(self, image_bytes: bytes) -> list[dict]:
|
|
"""返回 [{"bbox": {x,y,w,h}, "class_name", "confidence"}]"""
|
|
|
|
|
|
class MockDetector(Detector):
|
|
"""开发联调用：校验图片可解析后返回固定检测结果。"""
|
|
|
|
def __init__(self, class_name: str = "healthy", confidence: float = 0.95):
|
|
self.class_name = class_name
|
|
self.confidence = confidence
|
|
|
|
def detect(self, image_bytes: bytes) -> list[dict]:
|
|
try:
|
|
with Image.open(io.BytesIO(image_bytes)) as img:
|
|
img.verify()
|
|
except Exception as exc:
|
|
raise ValueError("无法解析图片") from exc
|
|
return [
|
|
{
|
|
"bbox": {"x": 10, "y": 10, "w": 100, "h": 100},
|
|
"class_name": self.class_name,
|
|
"confidence": self.confidence,
|
|
}
|
|
]
|
|
|
|
|
|
class ONNXDetector(Detector):
|
|
"""YOLOv8 导出的 best.onnx 推理（CPU）。
|
|
|
|
YOLO 训练恢复后放置模型文件，MODEL_MODE=onnx 生效。
|
|
输出张量按 YOLOv8 常见格式 [1, 4+nc, 8400]（或转置）解析，含 NMS。
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model_path: str,
|
|
labels: tuple[str, ...],
|
|
input_size: int = 640,
|
|
conf_threshold: float = 0.25,
|
|
iou_threshold: float = 0.45,
|
|
):
|
|
import numpy as np
|
|
import onnxruntime as ort
|
|
|
|
if not os.path.exists(model_path):
|
|
raise FileNotFoundError(f"模型文件不存在: {model_path}")
|
|
self.session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
|
|
self.labels = labels
|
|
self.input_size = input_size
|
|
self.conf_threshold = conf_threshold
|
|
self.iou_threshold = iou_threshold
|
|
self.np = np
|
|
|
|
def detect(self, image_bytes: bytes) -> list[dict]:
|
|
np = self.np
|
|
img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
|
|
orig_w, orig_h = img.size
|
|
scale = min(self.input_size / orig_w, self.input_size / orig_h)
|
|
new_w, new_h = int(round(orig_w * scale)), int(round(orig_h * scale))
|
|
resized = img.resize((new_w, new_h))
|
|
canvas = Image.new("RGB", (self.input_size, self.input_size), (114, 114, 114))
|
|
pad_x = (self.input_size - new_w) // 2
|
|
pad_y = (self.input_size - new_h) // 2
|
|
canvas.paste(resized, (pad_x, pad_y))
|
|
blob = np.asarray(canvas, dtype=np.float32) / 255.0
|
|
blob = blob.transpose(2, 0, 1)[None, ...]
|
|
out = self.session.run(None, {self.session.get_inputs()[0].name: blob})[0]
|
|
|
|
nc = len(self.labels)
|
|
if out.shape[1] != 4 + nc:
|
|
out = out.transpose(0, 2, 1)
|
|
out = out[0]
|
|
if out.shape[0] == 4 + nc:
|
|
out = out.T
|
|
|
|
boxes, scores, class_ids = [], [], []
|
|
for row in out:
|
|
cx, cy, w, h = row[:4]
|
|
cls_scores = row[4 : 4 + nc]
|
|
cls_id = int(cls_scores.argmax())
|
|
score = float(cls_scores[cls_id])
|
|
if score < self.conf_threshold:
|
|
continue
|
|
boxes.append(
|
|
[
|
|
(cx - w / 2 - pad_x) / scale,
|
|
(cy - h / 2 - pad_y) / scale,
|
|
w / scale,
|
|
h / scale,
|
|
]
|
|
)
|
|
scores.append(score)
|
|
class_ids.append(cls_id)
|
|
|
|
keep = self._nms(np.asarray(boxes), np.asarray(scores))
|
|
return [
|
|
{
|
|
"bbox": {
|
|
"x": float(boxes[i][0]),
|
|
"y": float(boxes[i][1]),
|
|
"w": float(boxes[i][2]),
|
|
"h": float(boxes[i][3]),
|
|
},
|
|
"class_name": self.labels[class_ids[i]],
|
|
"confidence": float(scores[i]),
|
|
}
|
|
for i in keep
|
|
]
|
|
|
|
def _nms(self, boxes, scores):
|
|
np = self.np
|
|
if len(boxes) == 0:
|
|
return []
|
|
x1 = boxes[:, 0]
|
|
y1 = boxes[:, 1]
|
|
x2 = boxes[:, 0] + boxes[:, 2]
|
|
y2 = boxes[:, 1] + boxes[:, 3]
|
|
areas = boxes[:, 2] * boxes[:, 3]
|
|
order = scores.argsort()[::-1]
|
|
keep = []
|
|
while order.size > 0:
|
|
i = int(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:]])
|
|
inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
|
|
iou = inter / (areas[i] + areas[order[1:]] - inter + 1e-6)
|
|
order = order[1:][iou <= self.iou_threshold]
|
|
return keep
|