orisys
419 lines
16 KiB
Python
419 lines
16 KiB
Python
"""工件二阶段检测：YOLO 定位工件 + HoughCircles 检测左右螺纹孔中心。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import cv2
|
||
import numpy as np
|
||
import yaml
|
||
|
||
ROOT = Path(__file__).resolve().parent
|
||
DEFAULT_CONFIG = ROOT / "hole_detect_config.yaml"
|
||
|
||
|
||
def load_config(path: str | Path | None = None) -> dict[str, Any]:
|
||
cfg_path = Path(path) if path else DEFAULT_CONFIG
|
||
with open(cfg_path, encoding="utf-8") as f:
|
||
data = yaml.safe_load(f) or {}
|
||
if not isinstance(data, dict):
|
||
raise ValueError(f"配置格式错误: {cfg_path}")
|
||
return data
|
||
|
||
|
||
def _fail(
|
||
msg: str,
|
||
draw_img: np.ndarray,
|
||
bbox: tuple[int, int, int, int] | None = None,
|
||
gate_block: bool = False,
|
||
) -> dict[str, Any]:
|
||
return {
|
||
"status": False,
|
||
"msg": msg,
|
||
"hole_points": [],
|
||
"hole_radii": [],
|
||
"bbox": list(bbox) if bbox is not None else None,
|
||
"gate_block": gate_block,
|
||
"draw_img": draw_img,
|
||
}
|
||
|
||
|
||
def _draw_bbox(img: np.ndarray, bbox: tuple[int, int, int, int], color=(0, 255, 0)) -> None:
|
||
x1, y1, x2, y2 = bbox
|
||
cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
|
||
cv2.putText(img, "gongjian", (x1, max(20, y1 - 8)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)
|
||
|
||
|
||
def _draw_center_gate(img: np.ndarray, max_dx: float, max_dy: float, ok: bool) -> None:
|
||
h, w = img.shape[:2]
|
||
cx, cy = w * 0.5, h * 0.5
|
||
x1 = int(round(cx - max_dx))
|
||
y1 = int(round(cy - max_dy))
|
||
x2 = int(round(cx + max_dx))
|
||
y2 = int(round(cy + max_dy))
|
||
color = (0, 255, 0) if ok else (0, 165, 255)
|
||
cv2.rectangle(img, (x1, y1), (x2, y2), color, 1)
|
||
cv2.drawMarker(img, (int(cx), int(cy)), color, markerType=cv2.MARKER_CROSS, markerSize=18, thickness=1)
|
||
|
||
|
||
def _bbox_center_ok(
|
||
bbox: tuple[int, int, int, int],
|
||
img_w: int,
|
||
img_h: int,
|
||
cfg: dict[str, Any],
|
||
) -> tuple[bool, float, float, float, float]:
|
||
"""返回 (是否居中, bbox_cx, bbox_cy, max_dx, max_dy)。"""
|
||
x1, y1, x2, y2 = bbox
|
||
bx = 0.5 * (x1 + x2)
|
||
by = 0.5 * (y1 + y2)
|
||
max_dx = float(cfg.get("center_max_offset_ratio_x", 0.12)) * img_w
|
||
max_dy = float(cfg.get("center_max_offset_ratio_y", 0.12)) * img_h
|
||
ok = abs(bx - img_w * 0.5) <= max_dx and abs(by - img_h * 0.5) <= max_dy
|
||
return ok, bx, by, max_dx, max_dy
|
||
|
||
|
||
def _draw_holes(
|
||
img: np.ndarray,
|
||
points: list[list[float]],
|
||
radii: list[float] | None = None,
|
||
) -> None:
|
||
"""画孔轮廓圆 + 中心点（十字+实心点）+ 坐标。"""
|
||
for i, (cx, cy) in enumerate(points):
|
||
px, py = int(round(cx)), int(round(cy))
|
||
r = 12.0 if radii is None or i >= len(radii) else float(radii[i])
|
||
rr = max(4, int(round(r)))
|
||
# 轮廓
|
||
cv2.circle(img, (px, py), rr, (0, 0, 255), 2)
|
||
# 中心十字 + 黄点
|
||
cv2.drawMarker(img, (px, py), (0, 0, 255), markerType=cv2.MARKER_CROSS, markerSize=16, thickness=2)
|
||
cv2.circle(img, (px, py), 3, (0, 255, 255), -1)
|
||
label = f"P{i + 1}({px},{py}) r={rr}"
|
||
cv2.putText(img, label, (px + 12, py - 8), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 0, 255), 2)
|
||
|
||
|
||
def _parse_best_bbox(dets, conf_thr: float) -> tuple[int, int, int, int] | None:
|
||
"""最多取 1 个工件：取置信度最高且 >= conf 的框。"""
|
||
if not dets:
|
||
return None
|
||
ordered = sorted(dets, key=lambda d: d.conf, reverse=True)
|
||
for det in ordered:
|
||
if float(det.conf) < conf_thr:
|
||
continue
|
||
return tuple(int(v) for v in det.xyxy) # type: ignore[return-value]
|
||
return None
|
||
|
||
|
||
def _pad_roi(
|
||
img_w: int,
|
||
img_h: int,
|
||
bbox: tuple[int, int, int, int],
|
||
padding: int,
|
||
) -> tuple[int, int, int, int, int, int]:
|
||
"""返回 x1,y1,x2,y2,offset_x,offset_y（裁剪左上角在原图坐标）。"""
|
||
x1, y1, x2, y2 = bbox
|
||
px1 = max(0, x1 - padding)
|
||
py1 = max(0, y1 - padding)
|
||
px2 = min(img_w, x2 + padding)
|
||
py2 = min(img_h, y2 + padding)
|
||
if px2 <= px1 or py2 <= py1:
|
||
raise ValueError("ROI 无效（宽或高为 0）")
|
||
return px1, py1, px2, py2, px1, py1
|
||
|
||
|
||
def _render_success(
|
||
image: np.ndarray,
|
||
bbox: tuple[int, int, int, int],
|
||
hole_points: list[list[float]],
|
||
hole_radii: list[float],
|
||
dist: float,
|
||
gate: tuple[float, float] | None = None,
|
||
) -> np.ndarray:
|
||
draw_img = image.copy()
|
||
if gate is not None:
|
||
_draw_center_gate(draw_img, gate[0], gate[1], True)
|
||
_draw_bbox(draw_img, bbox)
|
||
_draw_holes(draw_img, hole_points, hole_radii)
|
||
cv2.putText(
|
||
draw_img,
|
||
f"OK dist={dist:.1f}px",
|
||
(16, 36),
|
||
cv2.FONT_HERSHEY_SIMPLEX,
|
||
0.8,
|
||
(0, 255, 0),
|
||
2,
|
||
)
|
||
return draw_img
|
||
|
||
|
||
def detect_holes(
|
||
image: np.ndarray,
|
||
detector,
|
||
config: dict[str, Any] | None = None,
|
||
device: str | int = 0,
|
||
) -> dict[str, Any]:
|
||
"""
|
||
二阶段检测螺纹孔中心。
|
||
|
||
detector: 需提供 detect(image_bgr, conf=..., iou=..., imgsz=...) -> list[DetBox]
|
||
（YOLODetector / FasterRCNNDetector 均可）
|
||
|
||
返回:
|
||
status, msg, hole_points [[cx,cy],...], hole_radii [r1,r2], bbox, draw_img(BGR)
|
||
"""
|
||
try:
|
||
if image is None or not isinstance(image, np.ndarray) or image.size == 0:
|
||
return _fail("输入图像无效", np.zeros((480, 640, 3), dtype=np.uint8))
|
||
|
||
cfg = config if config is not None else load_config()
|
||
draw_img = image.copy()
|
||
h, w = image.shape[:2]
|
||
|
||
conf = float(cfg.get("conf", 0.25))
|
||
iou = float(cfg.get("iou", 0.7))
|
||
imgsz = int(cfg.get("imgsz", 640))
|
||
padding = int(cfg.get("bbox_padding", 25))
|
||
ksize = int(cfg.get("gaussian_ksize", 5))
|
||
if ksize % 2 == 0:
|
||
ksize += 1
|
||
sigma = float(cfg.get("gaussian_sigma", 1.5))
|
||
|
||
# 1) 检测器定位工件 bbox（最多 1 个）— 兼容 YOLO / Faster R-CNN
|
||
dets = detector.detect(image, conf=conf, iou=iou, imgsz=imgsz)
|
||
bbox = _parse_best_bbox(dets, conf)
|
||
if bbox is None:
|
||
return _fail("未检测到工件", draw_img)
|
||
|
||
_draw_bbox(draw_img, bbox)
|
||
|
||
# 1.5) 居中门控：bbox 中心接近画面中心才进入二阶检测
|
||
center_enable = bool(cfg.get("center_gate_enable", True))
|
||
centered, bx, by, max_dx, max_dy = _bbox_center_ok(bbox, w, h, cfg)
|
||
if bool(cfg.get("center_gate_draw", True)):
|
||
_draw_center_gate(draw_img, max_dx, max_dy, centered if center_enable else True)
|
||
if center_enable and not centered:
|
||
dx = abs(bx - w * 0.5)
|
||
dy = abs(by - h * 0.5)
|
||
cv2.putText(
|
||
draw_img,
|
||
"MOVE TO CENTER",
|
||
(16, 36),
|
||
cv2.FONT_HERSHEY_SIMPLEX,
|
||
0.8,
|
||
(0, 165, 255),
|
||
2,
|
||
)
|
||
return _fail(
|
||
f"工件未居中：offset=({dx:.0f},{dy:.0f})px，允许<=({max_dx:.0f},{max_dy:.0f})px",
|
||
draw_img,
|
||
bbox,
|
||
gate_block=True,
|
||
)
|
||
|
||
# 2) padding ROI + 边界保护
|
||
rx1, ry1, rx2, ry2, offset_x, offset_y = _pad_roi(w, h, bbox, padding)
|
||
roi = image[ry1:ry2, rx1:rx2]
|
||
if roi.size == 0:
|
||
return _fail("ROI 裁剪失败", draw_img, bbox)
|
||
|
||
# 3) 灰度 + 高斯模糊
|
||
gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY) if roi.ndim == 3 else roi
|
||
blur = cv2.GaussianBlur(gray, (ksize, ksize), sigma)
|
||
|
||
# 4) HoughCircles
|
||
circles = cv2.HoughCircles(
|
||
blur,
|
||
cv2.HOUGH_GRADIENT,
|
||
dp=float(cfg.get("hough_dp", 1.2)),
|
||
minDist=float(cfg.get("hough_min_dist", 40)),
|
||
param1=float(cfg.get("hough_param1", 100)),
|
||
param2=float(cfg.get("hough_param2", 28)),
|
||
minRadius=int(cfg.get("hough_min_radius", 6)),
|
||
maxRadius=int(cfg.get("hough_max_radius", 45)),
|
||
)
|
||
if circles is None:
|
||
return _fail("HoughCircles 未检测到圆孔", draw_img, bbox)
|
||
|
||
# 5) ROI 坐标 -> 原图坐标
|
||
circles = np.asarray(circles[0], dtype=np.float64)
|
||
mapped: list[tuple[float, float, float]] = []
|
||
for cx, cy, r in circles:
|
||
mapped.append((float(cx + offset_x), float(cy + offset_y), float(r)))
|
||
|
||
# 6) 工业校验
|
||
bx1, by1, bx2, by2 = bbox
|
||
min_r = float(cfg.get("min_radius", 6))
|
||
max_r = float(cfg.get("max_radius", 45))
|
||
dist_min = float(cfg.get("hole_distance_min", 80))
|
||
dist_max = float(cfg.get("hole_distance_max", 900))
|
||
|
||
valid: list[tuple[float, float, float]] = []
|
||
for cx, cy, r in mapped:
|
||
if not (bx1 <= cx <= bx2 and by1 <= cy <= by2):
|
||
continue
|
||
if not (min_r <= r <= max_r):
|
||
continue
|
||
valid.append((cx, cy, r))
|
||
|
||
# 多于 2 个时取最左/最右一对（对应左右螺纹孔），仍须满足间距规则
|
||
if len(valid) > 2:
|
||
ordered = sorted(valid, key=lambda c: c[0])
|
||
left, right = ordered[0], ordered[-1]
|
||
d = float(np.hypot(left[0] - right[0], left[1] - right[1]))
|
||
if dist_min <= d <= dist_max:
|
||
valid = [left, right]
|
||
else:
|
||
return _fail(
|
||
f"候选圆孔={len(valid)}，左右端点距离={d:.1f}px 不在 [{dist_min}, {dist_max}]",
|
||
draw_img,
|
||
bbox,
|
||
)
|
||
|
||
if len(valid) != 2:
|
||
return _fail(f"有效圆孔数量={len(valid)}，要求严格等于 2", draw_img, bbox)
|
||
|
||
ordered = sorted(valid, key=lambda c: c[0])
|
||
(x1, y1, r1), (x2, y2, r2) = ordered
|
||
dist = float(np.hypot(x1 - x2, y1 - y2))
|
||
if not (dist_min <= dist <= dist_max):
|
||
return _fail(
|
||
f"两孔距离={dist:.1f}px，不在 [{dist_min}, {dist_max}]",
|
||
draw_img,
|
||
bbox,
|
||
)
|
||
|
||
hole_points = [[x1, y1], [x2, y2]]
|
||
hole_radii = [r1, r2]
|
||
gate = (max_dx, max_dy) if bool(cfg.get("center_gate_draw", True)) else None
|
||
draw_img = _render_success(image, bbox, hole_points, hole_radii, dist, gate=gate)
|
||
return {
|
||
"status": True,
|
||
"msg": f"检测成功，两孔距离 {dist:.1f}px",
|
||
"hole_points": hole_points,
|
||
"hole_radii": hole_radii,
|
||
"bbox": list(bbox),
|
||
"draw_img": draw_img,
|
||
}
|
||
except Exception as exc:
|
||
fallback = image.copy() if isinstance(image, np.ndarray) and image.size else np.zeros((480, 640, 3), dtype=np.uint8)
|
||
return _fail(f"异常: {exc}", fallback)
|
||
|
||
|
||
class HoleStabilizer:
|
||
"""时序稳帧：EMA 平滑中心/半径，短暂失败时沿用上一帧结果，抑制闪烁。"""
|
||
|
||
def __init__(self, config: dict[str, Any] | None = None) -> None:
|
||
cfg = config or {}
|
||
self.cfg = cfg
|
||
self.hold_frames = int(cfg.get("stabilize_hold_frames", 15))
|
||
self.ema_alpha = float(cfg.get("stabilize_ema_alpha", 0.35))
|
||
self.max_jump = float(cfg.get("stabilize_max_jump", 60))
|
||
self.reset()
|
||
|
||
def reset(self) -> None:
|
||
self.points: list[list[float]] | None = None
|
||
self.radii: list[float] | None = None
|
||
self.bbox: tuple[int, int, int, int] | None = None
|
||
self.miss = 0
|
||
|
||
def update(self, image: np.ndarray, result: dict[str, Any]) -> dict[str, Any]:
|
||
try:
|
||
# 未居中：不做二阶、不沿用旧孔，立即清空稳帧
|
||
if result.get("gate_block"):
|
||
self.reset()
|
||
out = dict(result)
|
||
out["hole_points"] = []
|
||
out["hole_radii"] = []
|
||
return out
|
||
|
||
bbox = result.get("bbox")
|
||
if bbox is not None and len(bbox) == 4:
|
||
self.bbox = (int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3]))
|
||
|
||
if result.get("status") and len(result.get("hole_points") or []) == 2:
|
||
pts = [[float(p[0]), float(p[1])] for p in result["hole_points"]]
|
||
pts = sorted(pts, key=lambda p: p[0])
|
||
rad = result.get("hole_radii") or [12.0, 12.0]
|
||
rad = [float(rad[0]), float(rad[1])] if len(rad) >= 2 else [12.0, 12.0]
|
||
|
||
if self.points is not None:
|
||
jump = max(
|
||
float(np.hypot(pts[0][0] - self.points[0][0], pts[0][1] - self.points[0][1])),
|
||
float(np.hypot(pts[1][0] - self.points[1][0], pts[1][1] - self.points[1][1])),
|
||
)
|
||
if jump > self.max_jump:
|
||
# 跳变过大，当作本帧失败，沿用稳定值
|
||
self.miss += 1
|
||
else:
|
||
a = self.ema_alpha
|
||
self.points = [
|
||
[
|
||
(1 - a) * self.points[0][0] + a * pts[0][0],
|
||
(1 - a) * self.points[0][1] + a * pts[0][1],
|
||
],
|
||
[
|
||
(1 - a) * self.points[1][0] + a * pts[1][0],
|
||
(1 - a) * self.points[1][1] + a * pts[1][1],
|
||
],
|
||
]
|
||
self.radii = [
|
||
(1 - a) * self.radii[0] + a * rad[0],
|
||
(1 - a) * self.radii[1] + a * rad[1],
|
||
]
|
||
self.miss = 0
|
||
else:
|
||
self.points = pts
|
||
self.radii = rad
|
||
self.miss = 0
|
||
else:
|
||
self.miss += 1
|
||
if self.miss > self.hold_frames:
|
||
self.points = None
|
||
self.radii = None
|
||
|
||
if self.points is None or self.radii is None:
|
||
# 彻底丢失：只保留当前帧绘制（可能仅有 bbox）
|
||
out = dict(result)
|
||
out["hole_points"] = []
|
||
out["hole_radii"] = []
|
||
return out
|
||
|
||
dist = float(
|
||
np.hypot(
|
||
self.points[0][0] - self.points[1][0],
|
||
self.points[0][1] - self.points[1][1],
|
||
)
|
||
)
|
||
bbox_draw = self.bbox
|
||
if bbox_draw is None and result.get("bbox") is not None:
|
||
b = result["bbox"]
|
||
bbox_draw = (int(b[0]), int(b[1]), int(b[2]), int(b[3]))
|
||
if bbox_draw is None:
|
||
# 无框时仍画孔
|
||
draw_img = image.copy()
|
||
_draw_holes(draw_img, self.points, self.radii)
|
||
else:
|
||
gate = None
|
||
if bool(self.cfg.get("center_gate_draw", True)):
|
||
ih, iw = image.shape[:2]
|
||
max_dx = float(self.cfg.get("center_max_offset_ratio_x", 0.12)) * iw
|
||
max_dy = float(self.cfg.get("center_max_offset_ratio_y", 0.12)) * ih
|
||
gate = (max_dx, max_dy)
|
||
draw_img = _render_success(image, bbox_draw, self.points, self.radii, dist, gate=gate)
|
||
|
||
held = self.miss > 0
|
||
return {
|
||
"status": True,
|
||
"msg": (
|
||
f"检测成功，两孔距离 {dist:.1f}px"
|
||
+ (f"（稳帧 hold={self.miss}）" if held else "")
|
||
),
|
||
"hole_points": [[p[0], p[1]] for p in self.points],
|
||
"hole_radii": [float(self.radii[0]), float(self.radii[1])],
|
||
"bbox": list(bbox_draw) if bbox_draw is not None else result.get("bbox"),
|
||
"draw_img": draw_img,
|
||
}
|
||
except Exception as exc:
|
||
fallback = image.copy() if isinstance(image, np.ndarray) and image.size else np.zeros((480, 640, 3), dtype=np.uint8)
|
||
return _fail(f"稳帧异常: {exc}", fallback)
|