YOLOv5+GraspNet实战：如何用Python快速搭建机械臂抓取系统（附完整代码）

从“看见”到“抓取”用YOLOv5与GraspNet构建高精度机械臂视觉抓取系统想象一下你面前的工作台上散落着几个不同形状的零件一台机械臂需要从中准确地识别并抓取一个特定的螺丝。这听起来像是科幻电影里的场景但今天借助像YOLOv5和GraspNet这样的开源工具任何具备Python基础的开发者都能在自己的工作台上实现这套系统。这不仅仅是简单的“识别”和“移动”而是让机器理解物体的三维姿态并规划出最稳妥的抓取动作。对于从事机器人、自动化、智能制造甚至是创意交互装置开发的工程师和研究者来说掌握这套从视觉感知到物理执行的技术栈意味着能将想法快速转化为可运行的demo解决从实验室到产线的实际抓取难题。本文将从零开始带你走过环境搭建、核心代码实现、系统集成与调试的完整路径避开我当初踩过的那些坑最终让你手头的机械臂“眼明手快”。1. 项目基石环境配置与核心库解析在开始敲代码之前一个稳定、兼容的开发环境是成功的一半。不同于简单的Web开发视觉抓取项目涉及计算机视觉、深度学习、机器人学等多个领域的库版本冲突是新手最常见的“拦路虎”。我的建议是从一开始就使用虚拟环境进行隔离管理。首先我们使用Conda来创建和管理环境它能很好地处理复杂的非Python依赖如某些C编译的库。# 创建并激活一个名为‘grasp_system’的Python 3.8环境 conda create -n grasp_system python3.8 conda activate grasp_system接下来安装PyTorch。这里需要根据你是否有CUDA支持的NVIDIA显卡来选择版本。对于大多数实验和中等复杂度的模型使用GPU可以带来数十倍的加速。# 访问 https://pytorch.org/get-started/locally/ 获取最新安装命令 # 例如对于CUDA 11.3的Linux系统 pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113然后安装本项目的两大核心YOLOv5和GraspNet-baseline。注意我们这里使用社区维护的GraspNet基准实现而非原始论文中可能更复杂的代码库这更利于快速上手。# 安装YOLOv5 (Ultralytics版本) pip install ultralytics # 克隆并安装GraspNet-baseline git clone https://github.com/graspnet/graspnet-baseline.git cd graspnet-baseline pip install -e . cd ..此外我们还需要一些辅助库来处理图像、坐标转换和可视化pip install opencv-python pillow numpy matplotlib scipy transforms3d注意transforms3d这个库对于处理旋转矩阵、四元数、欧拉角之间的转换至关重要在机器人位姿处理中会频繁使用务必安装。为了让你对所需的核心库及其作用一目了然我整理了下表库名称主要用途在本项目中的角色安装备注PyTorch深度学习框架承载YOLOv5和GraspNet模型运行需匹配CUDA版本Ultralytics YOLOv5实时物体检测从图像中定位目标物体通过pip直接安装GraspNet-baseline6D抓取位姿估计预测物体在相机坐标系下的位置和姿态从GitHub克隆安装OpenCV计算机视觉图像读取、预处理、显示opencv-pythonNumPy数值计算矩阵运算、数据存储基础科学计算库transforms3d三维几何变换旋转、平移等位姿表示与转换处理机械臂坐标环境配好后不要急着跑代码。先用几行简单的命令验证关键库是否就绪import torch print(fPyTorch版本: {torch.__version__}) print(fCUDA是否可用: {torch.cuda.is_available()}) import cv2 print(fOpenCV版本: {cv2.__version__})如果一切正常你将看到PyTorch版本号和CUDA可用状态为True如果你安装了GPU版。至此你的“数字工具箱”已经准备妥当。2. 视觉感知层YOLOv5的实战调优与目标锁定YOLOv5以其出色的速度和精度平衡成为了工业界和学术界物体检测的首选之一。但在机械臂抓取场景中我们并非简单调用API就万事大吉。你需要考虑的是检测的稳定性、边界框的精度以及如何将2D图像坐标与3D世界联系起来。首先让我们加载一个预训练的YOLOv5模型并对一张图片进行检测。这里我推荐使用yolov5s模型它是“小”模型在保持不错精度的同时速度最快非常适合实时性要求高的抓取系统。import torch import cv2 from pathlib import Path # 加载模型 model torch.hub.load(ultralytics/yolov5, yolov5s, pretrainedTrue) # 切换到推理模式 model.eval() # 读取图像 img_path workbench_scene.jpg img cv2.imread(img_path) if img is None: raise FileNotFoundError(f无法在路径 {img_path} 找到图像文件) # 执行推理 results model(img) # 结果解析与可视化 results.print() # 在控制台打印检测到的类别、置信度和位置 results.show() # 弹出窗口显示带标注框的图像 # 获取结构化数据 detections results.pandas().xyxy[0] # 返回一个Pandas DataFrame print(detections.head())运行后你会看到一个弹窗图像中的物体被不同颜色的矩形框标出控制台会打印类似下面的信息image 1/1: 640x480 2 persons, 1 cup, 1 laptop Speed: 15.2ms pre-process, 45.6ms inference, 2.1ms NMS per image at shape (1, 3, 640, 480)DataFrame则包含了每个检测框的像素坐标、置信度和类别名称。然而在抓取系统中我们通常只关心特定类别的物体。假设我们要抓取的是“杯子”cup我们需要从所有检测结果中过滤出目标。# 定义目标类别 TARGET_CLASS cup # 过滤出目标类别的检测结果 target_detections detections[detections[name] TARGET_CLASS] if target_detections.empty: print(未检测到目标物体) else: # 通常选择置信度最高的那个检测框 best_detection target_detections.iloc[target_detections[confidence].argmax()] x1, y1, x2, y2 best_detection[[xmin, ymin, xmax, ymax]].astype(int).values confidence best_detection[confidence] print(f目标 {TARGET_CLASS} 检测到置信度: {confidence:.2f}) print(f边界框坐标: 左上({x1}, {y1}), 右下({x2}, {y2})) # 在图像上绘制目标框用绿色高亮 cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 3) cv2.putText(img, f{TARGET_CLASS}: {confidence:.2f}, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2) cv2.imshow(Target Detection, img) cv2.waitKey(0)这里有几个实战技巧置信度阈值results对象内部有一个默认置信度阈值如0.25。对于要求严格的工业场景你可以在推理时通过model.conf 0.6来调高阈值减少误检。多目标处理如果场景中有多个同类物体你需要设计策略来选择抓取哪一个比如选择离图像中心最近的或者尺寸最合适的。这可以通过计算边界框中心与图像中心的距离来实现。ROI裁剪得到2D边界框后我们通常需要将这个区域裁剪出来作为后续GraspNet的输入。这能帮助GraspNet专注于目标物体减少背景干扰。# 根据边界框裁剪目标区域 target_roi img[y1:y2, x1:x2] if target_roi.size 0: print(裁剪区域无效请检查坐标。) else: cv2.imwrite(target_cropped.jpg, target_roi) print(目标区域已保存为 target_cropped.jpg)至此你的系统已经拥有了“眼睛”能准确地从杂乱场景中找到并锁定那个需要被抓取的杯子。接下来我们要让系统理解这个杯子的三维姿态。3. 从2D到6DGraspNet深度解析与位姿估计实战物体检测告诉我们“目标在哪里2D框”而6D位姿估计则要回答“目标在三维空间中以何种姿态存在”。这包括3个平移自由度X, Y, Z坐标和3个旋转自由度围绕X, Y, Z轴的旋转。GraspNet这类深度学习模型通过学习大量物体与抓取姿态的对应关系能够从单张RGB图像或RGB-D图像中直接预测出可行的抓取位姿。首先我们需要加载预训练的GraspNet模型。由于GraspNet-baseline项目提供了预训练权重我们可以直接使用。确保你已经按照第一节的步骤克隆并安装了该库。import torch from graspnet_baseline.models import GraspNet from PIL import Image import torchvision.transforms as transforms # 初始化模型 device torch.device(cuda if torch.cuda.is_available() else cpu) model GraspNet().to(device) # 加载预训练权重 # 注意你需要从GraspNet-baseline的发布页面或指定路径下载权重文件 weights_path ./graspnet-baseline/checkpoints/your_graspnet_model.pth model.load_state_dict(torch.load(weights_path, map_locationdevice)) model.eval() print(GraspNet模型加载完毕运行在:, device)接下来准备输入图像。GraspNet通常需要固定尺寸的输入并且要进行归一化等预处理。def preprocess_for_graspnet(image_path, target_size224): 预处理图像以供GraspNet模型使用。 参数: image_path: 输入图像路径最好是上一步裁剪出的ROI target_size: 模型期望的输入尺寸 返回: image_tensor: 预处理后的图像张量 (1, C, H, W) # 使用PIL打开图像并转换为RGB image Image.open(image_path).convert(RGB) # 定义预处理流程 preprocess transforms.Compose([ transforms.Resize(target_size), transforms.CenterCrop(target_size), transforms.ToTensor(), # 使用ImageNet的均值和标准差进行归一化这是许多预训练模型的默认设置 transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) # 应用预处理并添加批次维度 image_tensor preprocess(image).unsqueeze(0) return image_tensor.to(device) # 使用上一节保存的裁剪图像 input_tensor preprocess_for_graspnet(target_cropped.jpg)现在进行6D位姿估计推理with torch.no_grad(): predictions model(input_tensor) # 解析预测结果 # 注意GraspNet的输出格式可能因版本和训练任务而异请查阅对应文档。 # 常见输出包括平移向量translation和旋转表示如四元数quaternion或旋转矩阵rotation_matrix translation predictions[translation].cpu().numpy()[0] # 形状 (3,) rotation predictions[rotation].cpu().numpy()[0] # 形状可能是 (4,) 四元数或 (3,3) 矩阵 print(预测的平移向量 (相机坐标系下):, translation) print(预测的旋转 (四元数 w, x, y, z):, rotation) # 假设输出是四元数到这里你可能会疑惑这个translation和rotation是相对于谁的在绝大多数情况下深度学习模型预测的位姿是相对于相机坐标系的。也就是说它假设相机是世界的原点预测出物体在相机前方的某个位置和朝向。这对于下一步——将位姿转换到机械臂坐标系——至关重要。为了更直观地理解GraspNet的输出我们来看一个典型的抓取位姿表示。一个完整的6D抓取位姿通常用一个4x4的齐次变换矩阵表示$$ T \begin{bmatrix} R t \ 0 1 \end{bmatrix} $$其中 $R$ 是一个3x3的旋转矩阵$t$ 是一个3x1的平移向量。GraspNet可能直接输出旋转矩阵也可能输出四元数更紧凑无奇异性。我们需要将其统一为变换矩阵以便后续坐标转换。import numpy as np from transforms3d.quaternions import quat2mat from transforms3d.affines import compose # 假设GraspNet输出的是四元数 [w, x, y, z] # 将四元数转换为旋转矩阵 rotation_matrix quat2mat(rotation) # 输入顺序为 [w, x, y, z] # 构建4x4齐次变换矩阵 grasp_pose_camera compose(translation, rotation_matrix, np.ones(3)) print(物体在相机坐标系下的抓取位姿矩阵:\n, grasp_pose_camera)这个grasp_pose_camera矩阵就是我们的核心成果之一。它描述了在相机“眼中”一个理想的抓取点例如机械手夹爪的中心应该处于什么位置和姿态。然而相机“眼中”的位置还需要转换到机械臂“手中”可执行的位置。4. 坐标之桥手眼标定与从相机到机械臂的坐标变换这是整个系统中最具工程挑战性的一环也是决定抓取精度的关键。我们得到了物体相对于相机的位姿grasp_pose_camera但机械臂控制器需要知道目标相对于其自身基座或末端工具的位姿。连接这两者的就是手眼标定Hand-Eye Calibration得到的变换矩阵。手眼标定解决的是“相机固定在哪里”的问题。有两种典型配置Eye-in-Hand相机安装在机械臂末端执行器上随机械臂移动。Eye-to-Hand相机固定在工作场景外部的某个位置静止不动。两者的标定方法略有不同但核心都是求解一个固定的变换矩阵相机坐标系到机械臂末端坐标系Eye-in-Hand或机械臂基座坐标系Eye-to-Hand的变换。这里我们以更常见的Eye-to-Hand相机固定为例。假设我们已经通过标定流程例如使用棋盘格和机械臂多次移动拍照得到了这个关键的变换矩阵T_camera_to_base。在实际项目中你可能使用OpenCV的calibrateHandEye函数或厂商提供的标定工具包来完成。# 假设通过标定得到的相机到机械臂基座的变换矩阵 # 这是一个4x4齐次变换矩阵 T_camera_to_base np.array([ [ 0.866, -0.500, 0.000, 0.500], [ 0.500, 0.866, 0.000, -0.200], [ 0.000, 0.000, 1.000, 1.000], [ 0.000, 0.000, 0.000, 1.000] ]) print(相机到机械臂基座的变换矩阵 T_camera_to_base:\n, T_camera_to_base)现在我们将相机坐标系下的抓取位姿转换到机械臂基座坐标系下# 计算物体在机械臂基座坐标系下的抓取位姿 # 公式: T_grasp_base T_camera_to_base * T_grasp_camera grasp_pose_base T_camera_to_base grasp_pose_camera print(物体在机械臂基座坐标系下的抓取位姿矩阵:\n, grasp_pose_base)得到grasp_pose_base后我们通常还需要考虑工具中心点Tool Center Point, TCP。机械臂末端的夹爪本身有尺寸我们需要告诉机械臂的是夹爪上某个特定点TCP应该到达的位置和姿态。因此可能还需要一个从夹爪法兰盘到TCP的固定变换T_tcp_to_flange。# 夹爪TCP相对于法兰盘的变换通常由夹爪厂商提供或通过测量得到 T_tcp_to_flange np.array([ [1, 0, 0, 0.00], [0, 1, 0, 0.00], [0, 0, 1, 0.05], # 假设TCP在法兰盘前方5cm处 [0, 0, 0, 1] ]) # 最终机械臂法兰盘需要到达的位姿是 # T_flange_base T_grasp_base * inv(T_tcp_to_flange) # 但更常见的做法是我们直接规划TCP的位姿即 grasp_pose_base。 # 许多机器人控制器允许直接设置TCP所以这里我们暂时使用 grasp_pose_base 作为目标。 target_pose_for_robot grasp_pose_base现在target_pose_for_robot这个4x4矩阵就是可以发送给机械臂控制器的终极指令。它精确描述了夹爪TCP应该到达的三维位置矩阵的第4列前三个元素和姿态矩阵的3x3旋转部分。5. 系统集成与抓取规划实战有了目标位姿下一步就是指挥机械臂行动。这里我们不会深入特定品牌机械臂如UR、ABB、Franka的SDK细节而是抽象出一个通用的控制接口并讨论抓取规划中的几个关键考量。首先我们需要从变换矩阵中提取出机械臂控制器通常需要的位姿表示形式位置X, Y, Z和欧拉角RX, RY, RZ或四元数。from transforms3d.euler import mat2euler def pose_matrix_to_robot_commands(pose_matrix): 将4x4齐次变换矩阵转换为机械臂常用的位置和欧拉角格式。 参数: pose_matrix: 4x4 numpy数组 返回: position: [x, y, z] 单位米 euler_angles: [rx, ry, rz] 单位弧度 (通常为ZYX欧拉角) # 提取平移部分 position pose_matrix[:3, 3] # 提取旋转矩阵并转换为ZYX欧拉角具体顺序需根据你的机器人控制器调整 rotation_matrix pose_matrix[:3, :3] # mat2euler 默认顺序是 sxyz即静态坐标系X, Y, Z轴旋转。 # 对于许多工业机器人使用的是ZYX欧拉角绕固定轴旋转。 euler_angles mat2euler(rotation_matrix, axessxyz) return position, euler_angles target_position, target_euler pose_matrix_to_robot_commands(target_pose_for_robot) print(f目标位置 (米): {target_position}) print(f目标欧拉角 (弧度): {target_euler}) print(f目标欧拉角 (度): {np.degrees(target_euler)})接下来我们模拟一个简单的抓取状态机。真实的抓取动作不是“瞬移”到目标点而是一系列安全、平滑的运动序列。class SimpleGraspPlanner: def __init__(self, robot_ipNone): # 这里可以初始化真实机器人的连接 # self.arm SomeRobotSDK(robot_ip) self.is_simulated True print(抓取规划器初始化完成模拟模式) def move_to_approach_pose(self, target_pose, offset_z0.1): 移动到接近点在目标点上方 offset_z 米处 approach_pose target_pose.copy() approach_pose[2, 3] offset_z # 在Z轴方向增加偏移 print(f[规划] 移动到接近点Z轴偏移 {offset_z}m) # self.arm.move_linear(approach_pose) return approach_pose def move_to_grasp_pose(self, target_pose, speed0.1): 直线运动到抓取点 print(f[规划] 直线运动到抓取点速度: {speed} m/s) # self.arm.move_linear(target_pose, speed) return True def execute_grasp(self): 执行抓取动作闭合夹爪 print([动作] 夹爪闭合) # self.arm.gripper_close() return True def move_to_retreat_pose(self, target_pose, offset_z0.15): 抬升到退回点 retreat_pose target_pose.copy() retreat_pose[2, 3] offset_z print(f[规划] 抬升到退回点Z轴偏移 {offset_z}m) # self.arm.move_linear(retreat_pose) return retreat_pose def run_grasp_sequence(self, target_pose_matrix): 运行完整的抓取序列 print( 开始执行抓取序列 ) # 1. 移动到接近点 self.move_to_approach_pose(target_pose_matrix) # 2. 下降到抓取点 self.move_to_grasp_pose(target_pose_matrix) # 3. 闭合夹爪 self.execute_grasp() # 4. 抬升到安全高度 self.move_to_retreat_pose(target_pose_matrix) # 5. 移动到放置点此处省略 # self.move_to_place_pose(...) print( 抓取序列执行完毕 ) # 实例化规划器并运行 planner SimpleGraspPlanner() planner.run_grasp_sequence(target_pose_for_robot)在实际部署中你还需要考虑更多细节碰撞检测在移动前需要检查路径是否与环境中其他物体发生碰撞。可以在运动规划中引入障碍物点云信息。抓取力控制对于易碎或形状不规则的物体需要控制夹爪的力度这可能需要力传感器反馈。误差补偿视觉误差、标定误差、机器人重复定位误差都会累积。工业中常采用“视觉伺服”或在最终抓取前进行二次精定位。最后将所有模块串联起来形成一个完整的、可循环执行的抓取系统主程序框架import time def main_loop(camera_index0, target_classcup): 主循环持续检测并抓取目标物体 # 初始化模块 detector torch.hub.load(ultralytics/yolov5, yolov5s) grasp_model GraspNet().to(device) grasp_model.load_state_dict(torch.load(weights_path)) grasp_model.eval() planner SimpleGraspPlanner() cap cv2.VideoCapture(camera_index) try: while True: ret, frame cap.read() if not ret: break # 步骤1: YOLOv5检测 results detector(frame) detections results.pandas().xyxy[0] target_det detections[detections[name] target_class] if not target_det.empty: best_det target_det.iloc[target_det[confidence].argmax()] x1, y1, x2, y2 best_det[[xmin, ymin, xmax, ymax]].astype(int).values # 步骤2: 裁剪ROI roi frame[y1:y2, x1:x2] if roi.size 0: continue cv2.imwrite(current_target.jpg, roi) # 步骤3: GraspNet 6D位姿估计 input_tensor preprocess_for_graspnet(current_target.jpg) with torch.no_grad(): pred grasp_model(input_tensor) trans, rot pred[translation].cpu().numpy()[0], pred[rotation].cpu().numpy()[0] grasp_pose_camera compose(trans, quat2mat(rot), np.ones(3)) # 步骤4: 坐标变换 grasp_pose_base T_camera_to_base grasp_pose_camera # 步骤5: 执行抓取 print(f发现目标 {target_class}开始抓取规划...) planner.run_grasp_sequence(grasp_pose_base) # 抓取后暂停一段时间或等待下一个指令 time.sleep(5) # 显示实时画面 cv2.imshow(Grasp System View, frame) if cv2.waitKey(1) 0xFF ord(q): break finally: cap.release() cv2.destroyAllWindows() # 运行主循环在实际使用中你可能需要先注释掉抓取执行部分仅测试视觉流程 # main_loop()在第一次运行完整系统时强烈建议采用“分步验证”和“模拟先行”的策略。先确保YOLOv5能稳定检测再验证GraspNet的位姿输出是否合理例如用手持物体在相机前移动观察预测位姿的变化是否符合预期最后再让机械臂空跑不带夹爪或远离物体运动轨迹确认坐标转换和运动规划是否正确。我自己的项目在集成阶段超过一半的时间都花在了这种循序渐进的调试和验证上但这能从根本上避免硬件损坏和安全隐患。当你看到机械臂第一次稳稳地抓起目标物体时那种成就感会让你觉得所有的调试都是值得的。