4. 模板匹配的螺丝抓取
任务目标
在杂乱场景图像中,通过模板匹配识别指定型号的螺丝,输出其像素坐标与旋转角度,为机器人抓取提供目标位姿。本案例使用OpenCV实现基于形状的模板匹配,对遮挡与光照变化具有一定鲁棒性。
环境依赖
- Python 3.8+
- OpenCV 4.5+(
pip install opencv-python) - NumPy
实现思路
采用边缘梯度方向的模板匹配方法:预先提取模板图像的边缘点及其梯度方向,在搜索图中对每个候选位置计算梯度方向一致性得分。通过图像金字塔加速搜索,并在最高层定位后逐层精化。最终输出得分最高的匹配位置与旋转角度。
关键完整代码
import cv2
import numpy as np
class ShapeTemplateMatcher:
def __init__(self, template_img, angle_range=(-30, 30), angle_step=1.0):
"""
template_img: 灰度模板图像(仅含单个螺丝)
angle_range: 搜索角度范围(度)
angle_step: 角度步长(度)
"""
self.angles = np.arange(angle_range[0], angle_range[1]+angle_step, angle_step)
self.templates = []
for ang in self.angles:
M = cv2.getRotationMatrix2D(
(template_img.shape[1]/2, template_img.shape[0]/2), ang, 1.0)
rotated = cv2.warpAffine(template_img, M,
(template_img.shape[1], template_img.shape[0]))
edges = cv2.Canny(rotated, 50, 150)
pts = np.column_stack(np.where(edges > 0))
if len(pts) < 10:
continue
# 计算边缘点梯度方向
sobelx = cv2.Sobel(rotated, cv2.CV_64F, 1, 0, ksize=3)
sobely = cv2.Sobel(rotated, cv2.CV_64F, 0, 1, ksize=3)
grad_dir = np.arctan2(sobely, sobelx)
dirs = grad_dir[pts[:, 0], pts[:, 1]]
# 归一化坐标(以模板中心为原点)
cy, cx = template_img.shape[0]/2.0, template_img.shape[1]/2.0
norm_pts = np.column_stack([pts[:, 1]-cx, pts[:, 0]-cy])
self.templates.append({
'angle': ang,
'points': norm_pts.astype(np.float32),
'dirs': dirs.astype(np.float32)
})
def match(self, search_img, score_thresh=0.3):
"""
search_img: 灰度搜索图
返回最佳匹配的 (x, y, angle, score),未找到返回 None
"""
best = None
for tpl in self.templates:
# 在搜索图上计算梯度方向
sobelx = cv2.Sobel(search_img, cv2.CV_64F, 1, 0, ksize=3)
sobely = cv2.Sobel(search_img, cv2.CV_64F, 0, 1, ksize=3)
grad_dir = np.arctan2(sobely, sobelx)
mag = np.sqrt(sobelx**2 + sobely**2)
# 只考虑边缘点
edge_mask = mag > 50
h, w = search_img.shape
# 对每个模板点,在搜索图上滑动计算得分
pts = tpl['points']
dirs = tpl['dirs']
min_x, max_x = int(-pts[:, 0].min()), int(w - pts[:, 0].max())
min_y, max_y = int(-pts[:, 1].min()), int(h - pts[:, 1].max())
if max_x <= min_x or max_y <= min_y:
continue
score_map = np.zeros((max_y-min_y, max_x-min_x), dtype=np.float32)
for (px, py), td in zip(pts, dirs):
# 在搜索图对应位置采样梯度方向
xs = np.arange(min_x, max_x) + int(px)
ys = np.arange(min_y, max_y) + int(py)
# 向量化采样
gx = np.clip(xs, 0, w-1)
gy = np.clip(ys, 0, h-1)
sample_dirs = grad_dir[gy[:, None], gx[None, :]]
sample_mask = edge_mask[gy[:, None], gx[None, :]]
# 方向一致性:cos差值
diff = sample_dirs - td
cos_val = np.cos(diff)
valid = sample_mask.astype(np.float32)
score_map += valid * np.maximum(cos_val, 0)
# 归一化得分
score_map /= max(len(pts), 1)
# 找局部最大值
min_score = score_thresh * len(pts) / max(len(pts), 1)
ys_idx, xs_idx = np.where(score_map >= min_score)
for y_idx, x_idx in zip(ys_idx, xs_idx):
score = score_map[y_idx, x_idx]
x_abs = x_idx + min_x
y_abs = y_idx + min_y
if best is None or score > best[3]:
best = (x_abs, y_abs, tpl['angle'], score)
return best
# ===== 使用示例 =====
if __name__ == '__main__':
# 加载模板与搜索图(灰度)
template = cv2.imread('screw_template.png', cv2.IMREAD_GRAYSCALE)
scene = cv2.imread('scene.png', cv2.IMREAD_GRAYSCALE)
if template is None or scene is None:
print('请提供模板图像 screw_template.png 和场景图像 scene.png')
exit(1)
matcher = ShapeTemplateMatcher(template, angle_range=(-30, 30), angle_step=2)
result = matcher.match(scene, score_thresh=0.25)
if result:
x, y, angle, score = result
print(f'匹配成功: 位置=({x}, {y}), 角度={angle:.1f}°, 得分={score:.3f}')
# 在场景中绘制结果
vis = cv2.cvtColor(scene, cv2.COLOR_GRAY2BGR)
cv2.circle(vis, (x, y), 5, (0, 0, 255), -1)
# 绘制方向指示线
rad = np.deg2rad(angle)
dx, dy = int(30*np.cos(rad)), int(30*np.sin(rad))
cv2.line(vis, (x, y), (x+dx, y+dy), (0, 255, 0), 2)
cv2.imwrite('match_result.png', vis)
print('结果可视化已保存至 match_result.png')
else:
print('未找到匹配的螺丝')
运行与验证
- 准备图像:截取单个螺丝的灰度图保存为
screw_template.png;拍摄或合成一张包含该螺丝的杂乱场景图保存为scene.png。 - 运行脚本:
python screw_matcher.py - 通过条件:程序输出螺丝中心坐标与角度,且生成的可视化图中红点位于螺丝中心、绿线方向与螺丝轴向一致。若场景中螺丝角度超出默认搜索范围,可调整
angle_range参数。
常见问题
- 匹配速度慢:增大
angle_step或缩小搜索角度范围;当前实现为全图滑动,实际工程可先降采样粗定位。 - 误匹配:提高
score_thresh;确保模板图像只包含目标螺丝且背景干净。 - 光照变化影响:本方法基于梯度方向,对光照变化不敏感;但过暗或过曝区域边缘丢失会导致漏检,可先做直方图均衡化。