AI计算机视觉笔记二十六:YOLOV8自训练关键点检测

简介: 本文档详细记录了使用YOLOv8训练关键点检测模型的过程。首先通过清华源安装YOLOv8,并验证安装。接着通过示例权重文件与测试图片`bus.jpg`演示预测流程。为准备训练数据,文档介绍了如何使用`labelme`标注工具进行关键点标注,并提供了一个Python脚本`labelme2yolo.py`将标注结果从JSON格式转换为YOLO所需的TXT格式。随后,通过Jupyter Notebook可视化标注结果确保准确性。最后,文档展示了如何组织数据集目录结构,并提供了训练与测试代码示例,包括配置文件`smoke.yaml`及训练脚本`train.py`,帮助读者完成自定义模型的训练与评估。

记录学习YOLOV8过程,自训练关键点检测模型。

清华源:-i https://mirror.baidu.com/pypi/simple

1、yolov8安装

git clone https://github.com/ultralytics/ultralytics
​
cd ultralytics
​
pip install -e .

安装成功后,使用命令 yolo 简单看下版本

(yolov8) llh@anhao:/$ yolo version
​
8.0.206

2、简单测试

下载权重文件

GitHub - ultralytics/ultralytics: NEW - YOLOv8 🚀 in PyTorch > ONNX > OpenVINO > CoreML > TFLite
image.png
直接点击下载。

获取测试图片的文件在ultralytics\assets目录有,使用的是的里面的bus.jpg测试。

使用 yolo 命令进行测试

yolo detect predict model=./yolov8n.pt source=./bus.jpg

image.png
输出在runs/detect/predict/目录下。
image.png

3、安装labelme

pip install labelme

直接在终端运行labelme打开软件
image.png
先取消“保存图片数据”(减少标注文件大小);在文件下
image.png
打开文件目录

image.png
接下来标注

先标注检测框,检测框用Create Rectangle(Ctrl+N)

填写类别名称

填写group_id,用于匹配后续标注的关键点,以当前画面中出现的顺序标注即可。
image.png
标注关键点,检测关键点用Create Point

按关键点顺序标注,如我们的顺序是head、tail,不可以错;

填写关键点名称,如这里是head;

填写关键点所在物体的group_id,匹配检测框

image.png
注意,如果多个类型需要填写group_id,group_id要匹配检测框

4、把JSON转成TXT

使用labelme标注生成的JSON文件,不能直接训练,所以需要把JSON文件转成TXT文件。

打开一个标注json文件,内容大致如下

version": "5.3.1",
  "flags": {},
  "shapes": [
    {
      "label": "smoke",
      "points": [
        [
          389.0,
          72.5
        ],
        [
          957.0,
          114.5
        ],
        [
          949.0,
          192.5
        ],
        [
          379.0,
          162.5
        ]
      ],
      "group_id": null,
      "description": "",
      "shape_type": "polygon",
      "flags": {}
    },

但是yolov8 要求的标注文件长这样

0 0.662641 0.494385 0.674719 0.988771 0.717187 0.189583 2.000000 0.798438 0.127083 2.000000 0.701562 0.091667 2.000000 0.921875 0.118750 2.000000 0.000000 0.000000 0.000000 0.971875 0.379167 2.000000 0.554688 0.262500 2.000000 0.000000 0.000000 0.000000 0.367188 0.427083 2.000000 0.767188 0.772917 2.000000 0.421875 0.500000 2.000000 0.829688 0.960417 1.000000 0.517188 0.881250 1.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000

根据这些规则,我们可以写一个转换脚本,将labelme标注的json格式转为yolo格式

labelme2yolo.py


# 将labelme标注的json文件转为yolo格式
import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
import glob
import json
import tqdm
​
​
# 物体类别
class_list = ["smoke"]
# 关键点的顺序
keypoint_list = ["head", "tail"]
​
​
def json_to_yolo(img_data,json_data):
    h,w = img_data.shape[:2]
    # 步骤:
    # 1. 找出所有的矩形,记录下矩形的坐标,以及对应group_id
    # 2. 遍历所有的head和tail,记下点的坐标,以及对应group_id,加入到对应的矩形中
    # 3. 转为yolo格式
​
    rectangles = {}
    # 遍历初始化
    for shape in json_data["shapes"]:
        label = shape["label"] # pen, head, tail
        group_id = shape["group_id"] # 0, 1, 2, ...
        points = shape["points"] # x,y coordinates
        shape_type = shape["shape_type"]
​
        # 只处理矩形
        if shape_type == "rectangle":
            if group_id not in rectangles:
                rectangles[group_id] = {
                    "label": label,
                    "rect": points[0] + points[1],  # Rectangle [x1, y1, x2, y2]
                    "keypoints_list": []
                }
    # 遍历更新,将点加入对应group_id的矩形中
    for keypoint in keypoint_list:
        for shape in json_data["shapes"]:
            label = shape["label"]
            group_id = shape["group_id"]
            points = shape["points"]
            # 如果匹配到了对应的keypoint
            if label == keypoint:
                rectangles[group_id]["keypoints_list"].append(points[0])

    # 转为yolo格式
    yolo_list = []
    for id, rectangle in rectangles.items():
        result_list  = []
        label_id = class_list.index(rectangle["label"])
        # x1,y1,x2,y2
        x1,y1,x2,y2 = rectangle["rect"]
        # center_x, center_y, width, height
        center_x = (x1+x2)/2
        center_y = (y1+y2)/2
        width = abs(x1-x2)
        height = abs(y1-y2)
        # normalize
        center_x /= w
        center_y /= h
        width /= w
        height /= h
​
        # 保留6位小数
        center_x = round(center_x, 6)
        center_y = round(center_y, 6)
        width = round(width, 6)
        height = round(height, 6)
​
​
        # 添加 label_id, center_x, center_y, width, height
        result_list = [label_id, center_x, center_y, width, height]
​
        # 添加 p1_x, p1_y, p1_v, p2_x, p2_y, p2_v
        for point in rectangle["keypoints_list"]:
            x,y = point
            x,y = int(x), int(y)
            # normalize
            x /= w
            y /= h
            # 保留6位小数
            x = round(x, 6)
            y = round(y, 6)

            result_list.extend([x,y,2])
​
        yolo_list.append(result_list)

    return yolo_list
​
​
# 获取所有的图片
img_list = glob.glob("./*.jpg")
​
for img_path in tqdm.tqdm( img_list ):

    img = cv2.imread(img_path)
    print(img_path)
    json_file = img_path.replace('jpg', 'json')
    with open(json_file) as json_file:
        json_data = json.load(json_file)
​
    yolo_list = json_to_yolo(img, json_data)

    yolo_txt_path = img_path.replace('jpg', 'txt')
    with open(yolo_txt_path, "w") as f:
        for yolo in yolo_list:
            for i in range(len(yolo)):
                if i == 0:
                    f.write(str(yolo[i]))
                else:
                    f.write(" " + str(yolo[i]))
            f.write("\n")
​

运行后,把JSON生成TXT。
image.png

5、对yolo格式的标注进行可视化。检查标注是否正确。

先安装jupyter

pip install jupyter
然后在终端输入

jupyter-lab


import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
import glob

img_path = './images/1.original_annotated/1 (20).jpg'

plt.figure(figsize=(15,10))
img = cv2.imread(img_path)
plt.imshow(img[:,:,::-1])
plt.axis('off')

yolo_txt_path = img_path.replace('jpg', 'txt')
print(yolo_txt_path)

with open(yolo_txt_path, 'r') as f:
    lines = f.readlines()

lines = [x.strip() for x in lines]

label = np.array([x.split() for x in lines], dtype=np.float32)

# 物体类别
class_list = ["smoke"]
# 类别的颜色
class_color = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255,255,0)]
# 关键点的顺序
keypoint_list = ["head", "tail"]
# 关键点的颜色
keypoint_color = [(255, 0, 0), (0, 255, 0)]

# 绘制检测框
img_copy = img.copy()
h,w = img_copy.shape[:2]
for id,l in enumerate( label ):
    # label_id ,center x,y and width, height
    label_id, cx, cy, bw, bh = l[0:5]
    label_text = class_list[int(label_id)]
    # rescale to image size
    cx *= w
    cy *= h
    bw *= w
    bh *= h

    # draw the bounding box
    xmin = int(cx - bw/2)
    ymin = int(cy - bh/2)
    xmax = int(cx + bw/2)
    ymax = int(cy + bh/2)
    cv2.rectangle(img_copy, (xmin, ymin), (xmax, ymax), class_color[int(label_id)], 2)
    cv2.putText(img_copy, label_text, (xmin, ymin-10), cv2.FONT_HERSHEY_SIMPLEX, 1, class_color[int(label_id)], 2)


# display the image
plt.figure(figsize=(15,10))
plt.imshow(img_copy[:,:,::-1])
plt.axis('off')
# save the image
cv2.imwrite("./tmp.jpg", img_copy)

img_copy = img.copy()
h,w = img_copy.shape[:2]
for id,l in enumerate( label ):
    # label_id ,center x,y and width, height
    label_id, cx, cy, bw, bh = l[0:5]
    label_text = class_list[int(label_id)]
    # rescale to image size
    cx *= w
    cy *= h
    bw *= w
    bh *= h

    # draw the bounding box
    xmin = int(cx - bw/2)
    ymin = int(cy - bh/2)
    xmax = int(cx + bw/2)
    ymax = int(cy + bh/2)
    cv2.rectangle(img_copy, (xmin, ymin), (xmax, ymax), class_color[int(label_id)], 2)
    cv2.putText(img_copy, label_text, (xmin, ymin-10), cv2.FONT_HERSHEY_SIMPLEX, 2, class_color[int(label_id)], 2)

    # draw 17 keypoints, px,py,pv,px,py,pv...
    for i in range(5, len(l), 3):
        px, py, pv = l[i:i+3]
        # rescale to image size
        px *= w
        py *= h
        # puttext the index 
        index = int((i-5)/3)
        # draw the keypoints
        cv2.circle(img_copy, (int(px), int(py)), 10, keypoint_color[int(index)], -1)
        keypoint_text = "{}_{}".format(index, keypoint_list[index])
        cv2.putText(img_copy, keypoint_text, (int(px), int(py)-10), cv2.FONT_HERSHEY_SIMPLEX, 1, keypoint_color[int(index)], 2)

plt.figure(figsize=(15,10))
plt.imshow(img_copy[:,:,::-1])
plt.axis('off')
# save 
cv2.imwrite('./tmp.jpg', img_copy)

image.png

6、训练

组织目录结构,我是直接把datasets目录放到yolov8的根目录下:
image.png
图片 datasets/custom_dataset/images/train/{文件名}.jpg对应的标注文件在 datasets/custom_dataset/labels/train/{文件名}.txt,YOLO会根据这个映射关系自动寻找(images换成labels);
训练集和验证集
images文件夹下有train和val文件夹,分别放置训练集和验证集图片;
labels文件夹有train和val文件夹,分别放置训练集和验证集标签(yolo格式)。
配置文件smoke.yaml

# Ultralytics YOLO 🚀, AGPL-3.0 license
# COCO8-pose dataset (first 8 images from COCO train2017) by Ultralytics
# Example usage: yolo train data=coco8-pose.yaml
# parent
# ├── ultralytics
# └── datasets
#     └── coco8-pose  ← downloads here (1 MB)


# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../datasets/custom_dataset  # dataset root dir
train: images/train  # train images (relative to 'path') 4 images
val: images/val  # val images (relative to 'path') 4 images
test:  # test images (optional)

# Keypoints
kpt_shape: [2, 3]  # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible)
#flip_idx: [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]

# Classes
names:
  0: smoke

# Download script/URL (optional)
#download: https://ultralytics.com/assets/coco8-pose.zip

編寫訓練train.py,訓練自己的模型

from ultralytics import YOLO

# 加载模型
# model = YOLO('yolov8s-pose.yaml')  # 从头训练
model = YOLO('./yolov8s-pose.pt')  # 使用预训练模型 (recommended for training)
# model = YOLO('yolov8s-pose.yaml').load('yolov8s-pose.pt')  # 从yaml构建网络并从预训练模型中迁移权重

# 训练
results = model.train(data='./smoke_pose.yaml', epochs=300, imgsz=640, workers=0, batch=20, project="pen_bolt", name="s120")

执行python train.py
image.png
会看到开始训练,并在pen_bolt\s1202下看到训练的结果。
image.png

7、测试

由于数据集才20张且只训练了120轮,效果不是太好,所以这里就直接使用图片测试

# 测试图片
from ultralytics import YOLO
import cv2
import numpy as np
import sys
# 读取命令行参数
weight_path = "./pen_bolt/s1202/weights/best.pt"
media_path = "./test.jpg"

# 加载模型
model = YOLO(weight_path )

# 获取类别
objs_labels = model.names  # get class labels
print(objs_labels)


# 类别的颜色
class_color = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255,255,0)]
# 关键点的顺序
keypoint_list = ["head", "tail"]
# 关键点的颜色
keypoint_color = [(255, 0, 0), (0, 255, 0)]

# 读取图片
frame = cv2.imread(media_path)
frame = cv2.resize(frame, (frame.shape[1]//2, frame.shape[0]//2))
# rotate
# 检测
result = list(model(frame, conf=0.3, stream=True))[0]  # inference,如果stream=False,返回的是一个列表,如果stream=True,返回的是一个生成器
boxes = result.boxes  # Boxes object for bbox outputs
boxes = boxes.cpu().numpy()  # convert to numpy array

# 遍历每个框
for box in boxes.data:
    l,t,r,b = box[:4].astype(np.int32) # left, top, right, bottom
    conf, id = box[4:] # confidence, class
    id = int(id)
    # 绘制框
    cv2.rectangle(frame, (l,t), (r,b), class_color[id], 2)
    # 绘制类别+置信度(格式:98.1%)
    cv2.putText(frame, f"{objs_labels[id]} {conf*100:.1f}%", (l, t-10), cv2.FONT_HERSHEY_SIMPLEX, 1, class_color[id], 2)

# 遍历keypoints
keypoints = result.keypoints  # Keypoints object for pose outputs
keypoints = keypoints.cpu().numpy()  # convert to numpy array

# draw keypoints, set first keypoint is red, second is blue
for keypoint in keypoints.data:
    for i in range(len(keypoint)):
        x,y,c = keypoint[i]
        x,y = int(x), int(y)
        cv2.circle(frame, (x,y), 10, keypoint_color[i], -1)
        cv2.putText(frame, f"{keypoint_list[i]}", (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 1, keypoint_color[i], 2)

    if len(keypoint) >= 2:
        # draw arrow line from tail to half between head and tail
        x1,y1,c1 = keypoint[0]
        x2,y2,c2 = keypoint[1]
        center_x, center_y = (x1+x2)/2, (y1+y2)/2
        cv2.arrowedLine(frame, (int(x2),int(y2)), (int(center_x), int(center_y)), (255,0,255), 4, line_type=cv2.LINE_AA, tipLength=0.1)


# save image
cv2.imwrite("result.jpg", frame)
print("save result.jpg")

直接python test.py测试,测试同一方向还好,旋转方向后有点不对,怀疑是数据集和参数问题,但整个过程是OK的。

image.png

相关文章
|
1天前
|
存储 人工智能 Cloud Native
小鹏汽车选用阿里云PolarDB,开启AI大模型训练新时代
PolarDB-PG云原生分布式数据库不仅提供了无限的扩展能力,还借助丰富的PostgreSQL生态系统,统一了后台技术栈,极大地简化了运维工作。这种强大的组合不仅提高了系统的稳定性和性能,还为小鹏汽车大模型训练的数据管理带来了前所未有的灵活性和效率。
|
7天前
|
机器学习/深度学习 人工智能 物联网
MiniMind:2小时训练出你的专属AI!开源轻量级语言模型,个人GPU轻松搞定
MiniMind 是一个开源的超小型语言模型项目,帮助开发者以极低成本从零开始训练自己的语言模型,最小版本仅需25.8M参数,适合在普通个人GPU上快速训练。
147 10
MiniMind:2小时训练出你的专属AI!开源轻量级语言模型,个人GPU轻松搞定
|
14天前
|
机器学习/深度学习 人工智能 Kubernetes
容器化AI模型部署实战:从训练到推理
在上一篇中,我们探讨了AI技术如何赋能容器化生态。本篇聚焦于AI模型的容器化部署,通过图像分类任务实例,详细介绍了从模型训练到推理服务的完整流程。使用PyTorch训练CNN模型,Docker打包镜像,并借助Kubernetes进行编排和部署,最终通过FastAPI提供推理服务。容器化技术极大提升了AI模型部署的便利性和管理效率,未来将成为主流趋势。
|
15天前
|
机器学习/深度学习 人工智能 PyTorch
从零开始:如何训练自己的AI模型
### 从零开始:如何训练自己的AI模型 训练AI模型如同培养新生儿,需耐心与技巧。首先明确目标(如图像识别、自然语言处理),选择框架(TensorFlow、PyTorch)。接着收集并预处理数据,确保多样性和准确性。然后设计模型结构,如卷积神经网络(CNN),并通过代码实现训练。训练后评估模型性能,调优以避免过拟合。最后部署模型至实际应用。通过猫狗分类器案例,掌握关键步骤和常见问题。训练AI模型是不断迭代优化的过程,实践才能真正掌握精髓。
|
17天前
|
人工智能 JSON PyTorch
TPO:告别微调!这个AI框架让大模型实时进化:无需训练直接优化,输入问题越用越聪明,输出质量暴涨50%
TPO(Test-Time Prompt Optimization)框架,通过奖励模型和迭代反馈优化大语言模型输出,无需训练即可显著提升性能,支持动态对齐人类偏好,降低优化成本。
174 8
TPO:告别微调!这个AI框架让大模型实时进化:无需训练直接优化,输入问题越用越聪明,输出质量暴涨50%
|
22天前
|
人工智能 Python
Light-A-Video:好莱坞级打光自由!上海AI Lab开源视频打光AI,无需训练秒改画面氛围,3步让阴天变夕阳
Light-A-Video 是由上海AI Lab联合交大等高校推出的无需训练的视频重照明方法,支持高质量、时间一致的光照控制,零样本生成和前景背景分离处理。
44 9
|
1月前
|
机器学习/深度学习 人工智能 计算机视觉
MILS:无需对LLM进行额外训练就能处理多模态任务,Meta AI提出零样本生成多模态描述方法
MILS 是 Meta AI 推出的零样本生成高质量多模态描述方法,支持图像、视频和音频的描述生成,无需额外训练。
134 34
MILS:无需对LLM进行额外训练就能处理多模态任务,Meta AI提出零样本生成多模态描述方法
|
6天前
|
存储 人工智能 自然语言处理
关于计算机视觉中的自回归模型,这篇综述一网打尽了
这篇综述文章全面介绍了自回归模型在计算机视觉领域的应用和发展。文章首先概述了视觉中的序列表示和建模基础知识,随后根据表示策略将视觉自回归模型分为基于像素、标记和尺度的三类框架,并探讨其与生成模型的关系。文章详细阐述了自回归模型在图像、视频、3D及多模态生成等多方面的应用,列举了约250篇参考文献,并讨论了其在新兴领域的潜力和面临的挑战。此外,文章还建立了一个GitHub存储库以整理相关论文,促进了学术合作与知识传播。论文链接:https://arxiv.org/abs/2411.05902
26 1
|
6月前
|
人工智能 测试技术 API
AI计算机视觉笔记二十 九:yolov10竹签模型,自动数竹签
本文介绍了如何在AutoDL平台上搭建YOLOv10环境并进行竹签检测与计数。首先从官网下载YOLOv10源码并创建虚拟环境,安装依赖库。接着通过官方模型测试环境是否正常工作。然后下载自定义数据集并配置`mycoco128.yaml`文件,使用`yolo detect train`命令或Python代码进行训练。最后,通过命令行或API调用测试训练结果,并展示竹签计数功能。如需转载,请注明原文出处。
|
6月前
|
人工智能 测试技术 PyTorch
AI计算机视觉笔记二十四:YOLOP 训练+测试+模型评估
本文介绍了通过正点原子的ATK-3568了解并实现YOLOP(You Only Look Once for Panoptic Driving Perception)的过程,包括训练、测试、转换为ONNX格式及在ONNX Runtime上的部署。YOLOP由华中科技大学团队于2021年发布,可在Jetson TX2上达到23FPS,实现了目标检测、可行驶区域分割和车道线检测的多任务学习。文章详细记录了环境搭建、训练数据准备、模型转换和测试等步骤,并解决了ONNX转换过程中的问题。

热门文章

最新文章