orisys/ros_flexiv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

first

Browse Source
This commit is contained in:
orisys
2026-08-21 14:51:57 +08:00
commit 4ad53f4e97
2056 changed files with 6272 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
hole_detect.py
+418
View File
@@ -0,0 +1,418 @@
"""工件二阶段检测：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)