基于YOLOv8的晶体管定位识别（pose），从数据集标注到训练手把手教程

原创

AI小怪兽

发布于 2024-01-22 09:26:33

2K40

代码可运行

文章被收录于专栏：YOLO大作战YOLO大作战

运行总次数：0

代码可运行

💡💡💡本文解决什么问题：教会你如何用晶体管从标注到训练Yolov8-pose关键点检测

1.如何标注自己的关键点数据集

1.1 labelme下载

# 安装labelme
pip install labelme

1.2使用labelme下

直接在python环境下运行

labelme

1.3 labelme介绍

关键点标记主要使用

1)Create Rectangle生成矩形框；

2）Create Point生成关键点；

1.4 数据集标注

2.数据集格式转换

2.1标记后的数据格式如下

一张图片对应一个json文件

json部分内容如下：

{
  "version": "5.1.1",
  "flags": {},
  "shapes": [
    {
      "label": "point",
      "points": [
        [
          300.29787234042556,
          156.6808510638298
        ]
      ],
      "group_id": null,
      "shape_type": "point",
      "flags": {}
    },
    {
      "label": "point1",
      "points": [
        [
          704.5531914893617,
          143.9148936170213
        ]
      ],
      "group_id": null,
      "shape_type": "point",
      "flags": {}
    },
    {
      "label": "point2",
      "points": [
        [
          716.2553191489362,
          551.3617021276597
        ]
      ],
      "group_id": null,
      "shape_type": "point",
      "flags": {}
    },
    {
      "label": "point3",
      "points": [
        [
          307.7446808510639,
          563.0638297872341
        ]
      ],
      "group_id": null,
      "shape_type": "point",
      "flags": {}
    },
    {
      "label": "transistor",
      "points": [
        [
          279.02127659574467,
          130.08510638297875
        ],
        [
          738.5957446808511,
          575.8297872340426
        ]
      ],
      "group_id": null,
      "shape_type": "rectangle",
      "flags": {}
    }
  ],

2.2 生成适合yolo格式的关键点数据集

labelme2yolo-keypoint

生成的txt内容如下:

0 0.49609 0.34375 0.44824 0.43457 0.29980 0.54980 2 0.29980 0.54980 2 0.29980 0.54980 2 0.29980 0.54980 2

讲解：

第一个0代表：框的类别，因为只有transistor一类，所以为0

0.49609 0.34375 0.44824 0.43457 代表：归一化后的框的中心点横纵坐标、宽、高

0.29980 0.54980 2 代表：归一化后的第一个关键点的横纵坐标、关键点可见性

关键点可见性理解：0代表不可见、1代表遮挡、2代表可见

2.3生成的yolo数据集如下

transistor：
-images:
 --train: png图片
 --val：png图片
-labels：
 --train: txt文件
 --val：txt文件

3.工业晶体管关键点训练

3.1 新建data/transistor/transistor.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: ./ultralytics-pose/data/transistor  # 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: [4, 3]  # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible)
flip_idx: [0, 1, 2, 3]

# Classes
names:
  0: transistor

3.2修改ultralytics/cfg/models/v8/yolov8-pose.yaml

修改为4个关键点和一个类别nc:1

# Ultralytics YOLO 🚀, AGPL-3.0 license
# YOLOv8-pose keypoints/pose estimation model. For Usage examples see https://docs.ultralytics.com/tasks/pose
 
# Parameters
nc: 1  # number of classes
kpt_shape: [21, 3]  # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible)
scales: # model compound scaling constants, i.e. 'model=yolov8n-pose.yaml' will call yolov8-pose.yaml with scale 'n'
  # [depth, width, max_channels]
  n: [0.33, 0.25, 1024]
  s: [0.33, 0.50, 1024]
  m: [0.67, 0.75, 768]
  l: [1.00, 1.00, 512]
  x: [1.00, 1.25, 512]


# YOLOv8.0n backbone
backbone:
  # [from, repeats, module, args]
  - [-1, 1, Conv, [64, 3, 2]]  # 0-P1/2
  - [-1, 1, Conv, [128, 3, 2]]  # 1-P2/4
  - [-1, 3, C2f, [128, True]]
  - [-1, 1, Conv, [256, 3, 2]]  # 3-P3/8
  - [-1, 6, C2f, [256, True]]
  - [-1, 1, Conv, [512, 3, 2]]  # 5-P4/16
  - [-1, 6, C2f, [512, True]]
  - [-1, 1, Conv, [1024, 3, 2]]  # 7-P5/32
  - [-1, 3, C2f, [1024, True]]
  - [-1, 1, SPPF, [1024, 5]]  # 9

# YOLOv8.0n head
head:
  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
  - [[-1, 6], 1, Concat, [1]]  # cat backbone P4
  - [-1, 3, C2f, [512]]  # 12

  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
  - [[-1, 4], 1, Concat, [1]]  # cat backbone P3
  - [-1, 3, C2f, [256]]  # 15 (P3/8-small)

  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 12], 1, Concat, [1]]  # cat head P4
  - [-1, 3, C2f, [512]]  # 18 (P4/16-medium)

  - [-1, 1, Conv, [512, 3, 2]]
  - [[-1, 9], 1, Concat, [1]]  # cat head P5
  - [-1, 3, C2f, [1024]]  # 21 (P5/32-large)

  - [[15, 18, 21], 1, Pose, [nc, kpt_shape]]  # Pose(P3, P4, P5)

3.3 默认参数开启训练

from ultralytics.cfg import entrypoint
arg="yolo pose train model=yolov8-pose.yaml data=data/transistor/transistor.yaml"
entrypoint(arg)

3.4训练结果分析

YOLOv8-pose summary (fused): 187 layers, 3077975 parameters, 0 gradients, 8.3 GFLOPs
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95)     Pose(P          R      mAP50  mAP50-95): 100%|██████████| 1/1 [00:00<00:00,  1.42it/s]
                   all          7          7      0.997      0.857      0.953      0.301      0.997      0.857      0.953      0.777
Speed: 0.4ms preprocess, 5.6ms inference, 0.0ms loss, 3.7ms postprocess per image