自用YOLO关键点检测之labelme2yolo格式改变

背景

labelme标注关键点，想要使用YOLO进行关键点检测时，需要将标签文件转换一下。

YOLO标签个格式

0 0.24830 0.34538 0.11451 0.15113 0.26254 0.41474 2 0.19686 0.30779 2 0.19977 0.32949 2
1

上面标签含义：
类型（0）
框位置（0.24830 0.34538 0.11451 0.15113）
关键点1的位置和是否可见（0.26254 0.41474 2）
关键点2的位置和是否可见（0.19686 0.30779 2）
关键点3的位置和是否可见（0.19977 0.32949 2）

0 无
1 遮挡
2 有
标注的时候应该是看见就标，看不见就不标

完整代码

bbox_class = {
    'Seed':0
}
keypoint_class = [
    'Seed_root', 'Seed_head1', 'Seed_head2', 'Seed_head3', 'Seed_head4',
    'Seed_head5', 'Seed_head6', 'Seed_head7', 'Seed_head8', 'Seed_head9',
    'Seed_head10', 'Seed_head11', 'Seed_head12', 'Seed_head13', 'Seed_head14',
    'Seed_head15', 'Seed_head16', 'Seed_head17'
]
# 设定文件夹路径
folder_path = r'keypoint1'
files = os.listdir(folder_path)

for file in files:
    if file.endswith('.json'):
        # 完整的文件路径
        labelme_path = os.path.join(folder_path, file)
        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.rsplit('.', 1)[0]
        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':
                    yolo_str = ''
                    bbox_class_id = bbox_class[each_ann['label']]
                    yolo_str += '{} '.format(bbox_class_id)
                    bbox = [min(each_ann['points'][0][0], each_ann['points'][1][0]),
                            max(each_ann['points'][0][0], each_ann['points'][1][0]),
                            min(each_ann['points'][0][1], each_ann['points'][1][1]),
                            max(each_ann['points'][0][1], each_ann['points'][1][1])]
                    bbox_center_x = (bbox[0] + bbox[1]) / 2
                    bbox_center_y = (bbox[2] + bbox[3]) / 2
                    bbox_width = bbox[1] - bbox[0]
                    bbox_height = bbox[3] - bbox[2]
                    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 = {}
                    for point_ann in labelme['shapes']:
                        if point_ann['shape_type'] == 'point':
                            x = point_ann['points'][0][0]
                            y = point_ann['points'][0][1]
                            label = point_ann['label']
                            if bbox[0] < x < bbox[1] and bbox[2] < y < bbox[3]:
                                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)
                        else:
                            yolo_str += '0 0 0 '
                    f.write(yolo_str + '\n')

        print('{} --> {} 转换完成'.format(labelme_path, yolo_txt_path))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67

修改位置

1.换成你的种类

bbox_class = {
    'Seed':0
}
1
2
3

2.按照你想要的关键点顺序换成你的关键点信息

keypoint_class = [
    'Seed_root', 'Seed_head1', 'Seed_head2', 'Seed_head3', 'Seed_head4',
    'Seed_head5', 'Seed_head6', 'Seed_head7', 'Seed_head8', 'Seed_head9',
    'Seed_head10', 'Seed_head11', 'Seed_head12', 'Seed_head13', 'Seed_head14',
    'Seed_head15', 'Seed_head16', 'Seed_head17'
]
1
2
3
4
5
6

3.更换成你的JSON文件的地址

folder_path = r'keypoint1'
1

你没标的会补上0 0 0，就是没有坐标也不可见的意思