YOLOV8 关键点检测之数据集制作

系统是ubuntu

软件：labelme  (labelImg不可以)

注意一点：左上角点击文件-------取消选中同时保存图像数据---------点击选中自动保存

为了方便寻找标注完的图像，可以自己创一个文件夹，然后修改输出路径

重点是如何把生成的json文件转为yolo需要的txt格式，我这里用的大佬的代码，很简单，两步到位：

1.先把json文件进行分类：分为训练集和验证集

代码：

import os
import shutil
 
from tqdm import tqdm
import random
#此代码为划分json格式的训练文件的
 
""" 使用：只需要修改 1. Dataset_folde，    
                  2. os.chdir(os.path.join(Dataset_folder, 'images'))里的 images,
                  3. val_scal = 0.2
                  4. os.chdir('../label_json')          label_json换成自己json标签文件夹名称   """
 
 
# 图片文件夹与json标签文件夹的根目录
Dataset_folder = r'/home/zdkj/桌面/111'
 
# 把当前工作目录改为指定路径
os.chdir(os.path.join(Dataset_folder, 'images'))   # images ： 图片文件夹的名称
folder = '.'                  # 代表os.chdir(os.path.join(Dataset_folder, 'images'))这个路径
imgs_list = os.listdir(folder)
 
random.seed(123)              # 固定随机种子，防止运行时出现bug后再次运行导致imgs_list 里面的图片名称顺序不一致
 
random.shuffle(imgs_list)     # 打乱
 
val_scal = 0.2                # 验证集比列
val_number = int(len(imgs_list) * val_scal)
val_files = imgs_list[:val_number]
train_files = imgs_list[val_number:]
 
print('all_files:', len(imgs_list))
print('train_files:', len(train_files))
print('val_files:', len(val_files))
 
 
os.mkdir('train')
for each in tqdm(train_files):
    shutil.move(each, 'train')
 
os.mkdir('val')
for each in tqdm(val_files):
    shutil.move(each, 'val')
 
os.chdir('../json')
 
os.mkdir('train')
for each in tqdm(train_files):
    json_file = os.path.splitext(each)[0] + '.json'
    shutil.move(json_file, 'train')
 
os.mkdir('val')
for each in tqdm(val_files):
    json_file = os.path.splitext(each)[0] + '.json'
    shutil.move(json_file, 'val')
 
print('划分完成')

 
 2.对分好类的json直接进行转换，一键转为txt

import os
import json
import shutil
import numpy as np
from tqdm import tqdm
 
""""#使用：1.改 bbox_class = {'sjb_rect': 0}，我的框的类别是sjb_rect，赋值它为0，如你是dog则改成：bbox_cls = {'dog': 0}
          2.改 keypoint_class = ['angle_30', 'angle_60', 'angle_90']，我的关键点类别是三个，分别是'angle_30', 'angle_60', 'angle_90'  
          3.改 Dataset_root 成你自己的图片与json文件的路径
          4.改 os.chdir('json_label/train')与os.chdir('json_label/val') 成你的json文件夹下的train与val文件夹 """
 
 
# 数据集根慕录(即图片文件夹与标签文件夹的上一级目录)
Dataset_root = '/home/zdkj/桌面/111'
 
# 框的类别
bbox_class = {
    'sensor':0
}
 
# 关键点的类别
keypoint_class = ['1', '2', '3', '4']   # 这里类别放的顺序对应关键点类别的标签 0，1，2
 
os.chdir(Dataset_root)
 
os.mkdir('labels')
os.mkdir('labels/train')
os.mkdir('labels/val')
 
 
def process_single_json(labelme_path, save_folder='../../labels/train'):
    with open(labelme_path, 'r', encoding='utf-8') as f:
        labelme = json.load(f)
 
    img_width = labelme['imageWidth']  # 图像宽度
    img_height = labelme['imageHeight']  # 图像高度
 
    # 生成 YOLO 格式的 txt 文件
    suffix = labelme_path.split('.')[-2]
    yolo_txt_path = suffix + '.txt'
 
    with open(yolo_txt_path, 'w', encoding='utf-8') as f:
 
        for each_ann in labelme['shapes']:  # 遍历每个标注
 
            if each_ann['shape_type'] == 'rectangle':  # 每个框，在 txt 里写一行
 
                yolo_str = ''
 
                ## 框的信息
                # 框的类别 ID
                bbox_class_id = bbox_class[each_ann['label']]
                yolo_str += '{} '.format(bbox_class_id)
                # 左上角和右下角的 XY 像素坐标
                bbox_top_left_x = int(min(each_ann['points'][0][0], each_ann['points'][1][0]))
                bbox_bottom_right_x = int(max(each_ann['points'][0][0], each_ann['points'][1][0]))
                bbox_top_left_y = int(min(each_ann['points'][0][1], each_ann['points'][1][1]))
                bbox_bottom_right_y = int(max(each_ann['points'][0][1], each_ann['points'][1][1]))
                # 框中心点的 XY 像素坐标
                bbox_center_x = int((bbox_top_left_x + bbox_bottom_right_x) / 2)
                bbox_center_y = int((bbox_top_left_y + bbox_bottom_right_y) / 2)
                # 框宽度
                bbox_width = bbox_bottom_right_x - bbox_top_left_x
                # 框高度
                bbox_height = bbox_bottom_right_y - bbox_top_left_y
                # 框中心点归一化坐标
                bbox_center_x_norm = bbox_center_x / img_width
                bbox_center_y_norm = bbox_center_y / img_height
                # 框归一化宽度
                bbox_width_norm = bbox_width / img_width
                # 框归一化高度
                bbox_height_norm = bbox_height / img_height
 
                yolo_str += '{:.5f} {:.5f} {:.5f} {:.5f} '.format(bbox_center_x_norm, bbox_center_y_norm,
                                                                  bbox_width_norm, bbox_height_norm)
 
                ## 找到该框中所有关键点，存在字典 bbox_keypoints_dict 中
                bbox_keypoints_dict = {}
                for each_ann in labelme['shapes']:  # 遍历所有标注
                    if each_ann['shape_type'] == 'point':  # 筛选出关键点标注
                        # 关键点XY坐标、类别
                        x = int(each_ann['points'][0][0])
                        y = int(each_ann['points'][0][1])
                        label = each_ann['label']
                        if (x > bbox_top_left_x) & (x < bbox_bottom_right_x) & (y < bbox_bottom_right_y) & (
                                y > bbox_top_left_y):  # 筛选出在该个体框中的关键点
                            bbox_keypoints_dict[label] = [x, y]
 
                ## 把关键点按顺序排好
                for each_class in keypoint_class:  # 遍历每一类关键点
                    if each_class in bbox_keypoints_dict:
                        keypoint_x_norm = bbox_keypoints_dict[each_class][0] / img_width
                        keypoint_y_norm = bbox_keypoints_dict[each_class][1] / img_height
                        yolo_str += '{:.5f} {:.5f} {} '.format(keypoint_x_norm, keypoint_y_norm,
                                                               2)  # 2-可见不遮挡 1-遮挡 0-没有点
                    else:  # 不存在的点，一律为0
                        yolo_str += '0 0 0 '
                # 写入 txt 文件中
                f.write(yolo_str + '\n')
 
    shutil.move(yolo_txt_path, save_folder)
    print('{} --> {} 转换完成'.format(labelme_path, yolo_txt_path))
 
 
 
os.chdir('json/train')
 
save_folder = '../../labels/train'
for labelme_path in os.listdir():
    try:
        process_single_json(labelme_path, save_folder=save_folder)
    except:
        print('******有误******', labelme_path)
print('YOLO格式的txt标注文件已保存至 ', save_folder)
 
 
os.chdir('../../')
 
os.chdir('json/val')
 
 
save_folder = '../../labels/val'
for labelme_path in os.listdir():
    try:
        process_single_json(labelme_path, save_folder=save_folder)
    except:
        print('******有误******', labelme_path)
print('YOLO格式的txt标注文件已保存至 ', save_folder)
 
os.chdir('../../')
 
 
os.chdir('../')

注意：第一个代码放在图片和json的文件夹外面，第二个代码和images、json放在一起。

最后，运行成功后会生成两个文件夹，train和labels，他们俩下面分别有两个子文件夹：train和val