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('未找到匹配的螺丝')

运行与验证

  1. 准备图像:截取单个螺丝的灰度图保存为screw_template.png;拍摄或合成一张包含该螺丝的杂乱场景图保存为scene.png。
  2. 运行脚本:python screw_matcher.py
  3. 通过条件:程序输出螺丝中心坐标与角度,且生成的可视化图中红点位于螺丝中心、绿线方向与螺丝轴向一致。若场景中螺丝角度超出默认搜索范围,可调整angle_range参数。

常见问题

  • 匹配速度慢:增大angle_step或缩小搜索角度范围;当前实现为全图滑动,实际工程可先降采样粗定位。
  • 误匹配:提高score_thresh;确保模板图像只包含目标螺丝且背景干净。
  • 光照变化影响:本方法基于梯度方向,对光照变化不敏感;但过暗或过曝区域边缘丢失会导致漏检,可先做直方图均衡化。