“YOLO26-obb+语音控制” 机械臂抓取方案（二、数据标注+模型训练）

m0_56498637

426人浏览 · 2026-04-01 16:49:24

m0_56498637 · 2026-04-01 16:49:24 发布

好的，我们在上午提到了数据的采集，那么我们现在开始做数据标注的工作，标注完成后我们就可以开始进行模型训练了。

数据标注

我们采集完的数据均放在images/train文件夹下，等待标注完成，我们再划分数据集（以8：2的比例），由于我们需要使用yolo26-obb，我们选择CVAT标注工具。

CVAT标注

CVAT标注工具在线网页：https://app.cvat.ai/tasks

有很多同学可能没有接触过这个工具，当然我也是第一次使用，之前都是用Labelimg或者Labelme等工具，我尝试了一下觉得也挺方便的，大家一起跟着来试一下吧。

标注指南

打开网页我们可以看到右上角+号创建新任务

创建任务界面

在这里我们将自己训练集图片都导入进来，task的Name可以随便取不影响

任务创建完成后，我们可以看到Add label选项，我们可以根据图片内你想识别的物体类别进行添加，添加完成后，我们标注的准备工作就算完成了

开始进入标注，点击Jobs一栏的 Job #你的任务编号，进入标注界面

这里可以使用快捷键N来进行画框，就和labelimg使用w画矩形框一样

标注完成后，左上角点击Menu，选择Upload annotations

导出后可以根据跳出的只是click here进行跳转，也可以回到主页点击request进行查看

随后我们就可以进行Download，下载导出的xml文件，需要进行解压，打开后的xml文件内容大概为下图所示

我们可以看到你的label名臣，包括一些矩形框信息以及旋转角度的信息

那么好，我们的标注工作就算彻底结束了，但是我们不能直接使用这个xml文件，因为我们需要转换成yolo格式，让每一个label映射到每一张图像上进行对应，所以这里我们写一个xml转yolo的py文件

xml2yolo.py

在这里你只需要修改文件的路径以及类别映射中你的label名称

import os
import math
import xml.etree.ElementTree as ET

# =========================
# 路径配置（自己改）
# =========================
XML_PATH = "../annotations_val.xml"
IMG_DIR = "../dataset/images/val"
LABEL_DIR = "../dataset/labels/val"

os.makedirs(LABEL_DIR, exist_ok=True)

# =========================
# 类别映射
# =========================
class_map = {
    "luckin_yellow": 0,
    "luckin_blue": 1,
}

# =========================
# 计算旋转框四点
# =========================
def get_obb_points(cx, cy, w, h, angle, img_w, img_h):
    angle_rad = math.radians(angle)

    cos_a = math.cos(angle_rad)
    sin_a = math.sin(angle_rad)

    dx = w / 2
    dy = h / 2

    # 四个角（顺时针）
    points = [
        (-dx, -dy),
        ( dx, -dy),
        ( dx,  dy),
        (-dx,  dy),
    ]

    result = []
    for x, y in points:
        rx = x * cos_a - y * sin_a
        ry = x * sin_a + y * cos_a

        px = (cx + rx) / img_w
        py = (cy + ry) / img_h

        result.append((px, py))

    return result


# =========================
# 解析XML
# =========================
tree = ET.parse(XML_PATH)
root = tree.getroot()

for image in root.findall("image"):
    file_name = image.get("name")
    width = float(image.get("width"))
    height = float(image.get("height"))

    label_path = os.path.join(
        LABEL_DIR,
        os.path.splitext(file_name)[0] + ".txt"
    )

    with open(label_path, "w") as f:
        for box in image.findall("box"):

            label = box.get("label")

            # 跳过未知类别
            if label not in class_map:
                print(f"⚠️ 未知类别: {label}")
                continue

            cls_id = class_map[label]

            xtl = float(box.get("xtl"))
            ytl = float(box.get("ytl"))
            xbr = float(box.get("xbr"))
            ybr = float(box.get("ybr"))

            # 中心点 + 宽高
            cx = (xtl + xbr) / 2
            cy = (ytl + ybr) / 2
            w = (xbr - xtl)
            h = (ybr - ytl)

            # rotation（可能没有）
            rotation_attr = box.get("rotation")
            angle = float(rotation_attr) if rotation_attr else 0.0

            # =========================
            # 生成四点
            # =========================
            points = get_obb_points(cx, cy, w, h, angle, width, height)

            # =========================
            # 写入（保留6位小数）
            # =========================
            line = f"{cls_id} " + " ".join(
                [f"{x:.6f} {y:.6f}" for x, y in points]
            ) + "\n"

            f.write(line)

    print(f"✅ 已生成: {label_path}")

print("🎉 全部转换完成！")

当控制台输出下面结果时，文件已经转换成功，可以对应点进去查看结果，这里就不给大家演示了

大家转换完txt文件后，不要忘记8:2的比例划分数据集，将图片和txt文件20%的数据放入val文件夹下保存

这里可以多做一次对于yolo格式的检查

data_check.py

import os

for file in os.listdir("../dataset/labels/train"):
    with open(os.path.join("../dataset/labels/train", file)) as f:
        for line in f:
            if len(line.strip().split()) != 6:
                print("错误:", file)
            else:
                print("yolo格式正常")

模型训练

我们已经完成了数据的标注，接下来我们将会使用yolo26-obb来进行训练

模型选择

我们从Ultralytics文档官网中下载obb模型，我这里下载的是yolo26n-obb.pt，现在模型都很轻量，下载都很快，就不提供网盘下载了。

官网链接：https://docs.ultralytics.com/zh/tasks/obb/#models

环境配置

在模型下载的同时我们也可以开始准备配置环境了，环境的配置可以参考https://blog.csdn.net/m0_56498637/article/details/159646240?spm=1001.2014.3001.5501

开始训练

Train.py

from ultralytics import YOLO

def main():
    model = YOLO("yolo26n-obb.pt")

    model.train(
        data=r"C:\Users\16095\PycharmProjects\Yolov11-obb+Vosk\data\yolo26-obb.yaml",
        epochs=100,
        imgsz=640,
        batch=8,
        device=0
    )

if __name__ == '__main__':
    main()

训练过程

当出现这一步的时候就代表已经满足训练要求，已经开始训练了

训练时，我们主要查看的值就是各类损失是否降低，查看mAP50和mAP50-95的值，我们不讲究的话，一般情况下这两个值基本上就对应了我们模型的识别准确率

训练后的结果我们可以在runs文件夹下查看

测试模型

图片测试

test_model_image.py

from ultralytics import YOLO
import cv2
import numpy as np
import math

MODEL_PATH = "../runs/obb/train/weights/best.pt"

names = {
    0: "luckin_yellow",
    1: "luckin_blue"
}

model = YOLO(MODEL_PATH)

img = cv2.imread("../test_images/1775027354.png")

results = model(img)

# =========================
# 核心函数：带方向的抓取角度
# =========================
def get_grasp_angle_with_direction(pts):
    """
    pts: (4,2) numpy array
    return: angle in degrees (-90 ~ 90)
    """

    edges = []

    for i in range(4):
        p1 = pts[i]
        p2 = pts[(i + 1) % 4]

        dx = p2[0] - p1[0]
        dy = p2[1] - p1[1]

        length = (dx ** 2 + dy ** 2) ** 0.5

        edges.append((length, dx, dy))

    # 找最长边
    _, dx, dy = max(edges, key=lambda x: x[0])

    # 统一方向（让dx为正）
    if dx < 0:
        dx = -dx
        dy = -dy

    angle = math.degrees(math.atan2(dy, dx))

    return angle

for r in results:
    if r.obb is None:
        continue

    boxes = r.obb.xyxyxyxy.cpu().numpy()
    cls_ids = r.obb.cls.cpu().numpy()
    scores = r.obb.conf.cpu().numpy()

    for box, cls_id, score in zip(boxes, cls_ids, scores):
        cls_id = int(cls_id)

        pts = box.reshape((4, 2))
        pts_int = pts.astype(int)

        # =========================
        # 计算角度
        # =========================
        p1, p2 = pts[0], pts[1]
        dx = p2[0] - p1[0]
        dy = p2[1] - p1[1]

        angle = get_grasp_angle_with_direction(pts)

        # =========================
        # 画框
        # =========================
        cv2.polylines(img, [pts_int], True, (0, 255, 0), 1)

        # =========================
        # 文本（类别+置信度+角度）
        # =========================
        label = f"{names[cls_id]} {score:.2f} {angle:.1f}deg"

        x, y = pts_int[0]

        cv2.putText(img, label, (x, y),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.5,
                    (0, 255, 0),
                    1,
                    cv2.LINE_AA)

cv2.imshow("result", img)
cv2.waitKey(0)
cv2.destroyAllWindows()

视频流检测

test_model_video.py

from ultralytics import YOLO
import pyrealsense2 as rs
import numpy as np
import cv2
import math

# =========================
# 模型路径
# =========================
MODEL_PATH = "../runs/obb/train/weights/best.pt"

names = {
    0: "luckin_yellow",
    1: "luckin_blue"
}

model = YOLO(MODEL_PATH)

# =========================
# RealSense 初始化
# =========================
pipeline = rs.pipeline()
config = rs.config()

config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
pipeline.start(config)

print("✅ RealSense 已启动")


# =========================
# 核心函数：带方向的抓取角度
# =========================
def get_grasp_angle_with_direction(pts):
    """
    pts: (4,2) numpy array
    return: angle in degrees (-90 ~ 90)
    """

    edges = []

    for i in range(4):
        p1 = pts[i]
        p2 = pts[(i + 1) % 4]

        dx = p2[0] - p1[0]
        dy = p2[1] - p1[1]

        length = (dx ** 2 + dy ** 2) ** 0.5

        edges.append((length, dx, dy))

    # 找最长边
    _, dx, dy = max(edges, key=lambda x: x[0])

    # 🔥 关键：统一方向（让dx为正）
    if dx < 0:
        dx = -dx
        dy = -dy

    angle = math.degrees(math.atan2(dy, dx))

    return angle


# =========================
# 主循环
# =========================
try:
    while True:
        frames = pipeline.wait_for_frames()
        color_frame = frames.get_color_frame()

        if not color_frame:
            continue

        frame = np.asanyarray(color_frame.get_data())

        # YOLO推理
        results = model(frame, verbose=False)

        for r in results:
            if r.obb is None:
                continue

            boxes = r.obb.xyxyxyxy.cpu().numpy()
            cls_ids = r.obb.cls.cpu().numpy()
            scores = r.obb.conf.cpu().numpy()

            for box, cls_id, score in zip(boxes, cls_ids, scores):
                cls_id = int(cls_id)

                pts = box.reshape((4, 2))
                pts_int = pts.astype(int)

                # =========================
                # 计算抓取角度
                # =========================
                angle = get_grasp_angle_with_direction(pts)
                if 0 <= angle <= 90:
                    angle = angle - 90
                elif -90 < angle < 0:
                    angle = angle + 90
                # =========================
                # 画OBB框
                # =========================
                cv2.polylines(frame, [pts_int], True, (0, 255, 0), 2)

                # =========================
                # 文本信息
                # =========================
                label = f"{names[cls_id]} {score:.2f} {angle:+.1f}deg"

                x, y = pts_int[0]

                cv2.putText(frame, label, (x, y),
                            cv2.FONT_HERSHEY_SIMPLEX,
                            0.5,  # 字体小一点
                            (0, 255, 0),
                            1,  # 更细
                            cv2.LINE_AA)

        cv2.imshow("RealSense OBB Detection", frame)

        if cv2.waitKey(1) & 0xFF == 27:
            break

finally:
    pipeline.stop()
    cv2.destroyAllWindows()

模型验证结果