yolov8旋转框+关键点检测

一、Yolov8obb_kpt

-----------------------------------现已在v8官方库上更新旋转框分割算法和旋转框关键点检测算法--------------------------
------------------------------------------- https://github.com/yzqxy/ultralytics-obb_segment---------------------------------
记得给博主点上⭐，给予博主继续创作好用算法的动力。
可参考博主上一篇文章，Yolov8obb_kpt，旋转框+关键点检测，有向目标检测，判断目标正方向

二、标注工具

标注软件：X-AnyLabeling，可查阅博客进行安装使用
标注旋转框和关键点，生成json文件。

三、json2txt转化为yolov8可训练标签的脚本

# COCO 格式的数据集转化为 YOLO 格式的数据集
# --json_path 输入的json文件路径
# --save_path 保存的文件夹名字，默认为当前目录下的labels。

import os
import json
from tqdm import tqdm
import argparse
import cv2
import numpy as np
parser = argparse.ArgumentParser()
# 这里根据自己的json文件位置，换成自己的就行
parser.add_argument('--json_path',
                    default=r'D:\Dataset\关键点\tading\coco_kpt_format\annotations\tading_kpt_val.json', type=str,
                    help="input: coco format(json)")
# 这里设置.txt文件保存位置
parser.add_argument('--save_path', default=r'D:\Dataset\关键点\tading\20220912_hailei\txt\val', type=str,
                    help="specify where to save the output dir of labels")
arg = parser.parse_args()


def convert(size, box):
    dw = 1. / (size[1])
    dh = 1. / (size[0])
    x = (box[0] + box[2]) / 2.0
    y = (box[1] + box[3] )/ 2.0
    w = box[2]-box[0]
    h = box[3]-box[1]

    x = round(x * dw, 6)
    w = round(w * dw, 6)
    y = round(y * dh, 6)
    h = round(h * dh, 6)
    return (x, y, w, h)


def check_points_in_rotated_boxes(points, boxes):
    """Check whether point is in rotated boxes

    Args:
        points (tensor): (1, L, 2) anchor points
        boxes (tensor): [B, N, 5] gt_bboxes
        eps (float): default 1e-9

    Returns:
        is_in_box (tensor): (B, N, L)

    """
    a = np.array(boxes[0])
    b = np.array(boxes[1])
    c = np.array(boxes[2])
    d = np.array(boxes[3])
    ab = b - a
    ad = d - a
    # [B, N, L, 2]
    ap = points - a
    # [B, N, L]
    norm_ab = np.sum(ab * ab)
    # [B, N, L]
    norm_ad = np.sum(ad * ad)
    # [B, N, L] dot product
    ap_dot_ab = np.sum(ap * ab)
    # [B, N, L] dot product
    ap_dot_ad = np.sum(ap * ad)
    # [B, N, L] <A, B> = |A|*|B|*cos(theta)
    is_in_box = (ap_dot_ab >= 0) & (ap_dot_ab <= norm_ab) & (ap_dot_ad >= 0) & (
            ap_dot_ad <= norm_ad)
    return is_in_box

def save_txt(data,ana_txt_save_path,obb_cls_kpt,obb_cls,width,heigth):
    with open(os.path.join(ana_txt_save_path, filename.split('.json')[0] + '.txt'), 'w', encoding='utf-8') as f:
        # 类别，obbbox，8个kpt
        for i in range(len(data['shapes'])):
            len_line = 9 + max(len(sublist) for sublist in obb_cls_kpt) * 3
            # len_line=9
            line = ['0' for _ in range(len_line)]
            for obb_i in range(len(obb_cls)):
                if data['shapes'][i]['label'] == obb_cls[obb_i]:
                    line[0] = str(obb_i)
                    print('obb_cls[obb_i]', obb_cls[obb_i])
                    # import  pdb
                    # pdb.set_trace()
                    x1 = data['shapes'][i]['points'][0][0]
                    y1 = data['shapes'][i]['points'][0][1]
                    x2 = data['shapes'][i]['points'][1][0]
                    y2 = data['shapes'][i]['points'][1][1]
                    x3 = data['shapes'][i]['points'][2][0]
                    y3 = data['shapes'][i]['points'][2][1]
                    x4 = data['shapes'][i]['points'][3][0]
                    y4 = data['shapes'][i]['points'][3][1]
                    # if x1>x2:
                    #     x3=x2
                    #     x2=x1
                    #     x1=x3
                    # if y1 > y2:
                    #     y3=y2
                    #     y2=y1
                    #     y1=y3

                    tatou = [[x1, y1], [x2, y2], [x3, y3], [x4, y4]]

                    line[1] = str(round(x1 / width, 6))
                    line[2] = str(round(y1 / heigth, 6))
                    line[3] = str(round(x2 / width, 6))
                    line[4] = str(round(y2 / heigth, 6))
                    line[5] = str(round(x3 / width, 6))
                    line[6] = str(round(y3 / heigth, 6))
                    line[7] = str(round(x4 / width, 6))
                    line[8] = str(round(y4 / heigth, 6))
                    print(line)
                    for j in range(len(data['shapes'])):
                        if len(obb_cls_kpt[obb_i]) > 0:
                            for kpt_i in range(len(obb_cls_kpt[obb_i])):
                                if data['shapes'][j]['label'] == obb_cls_kpt[obb_i][kpt_i]:
                                    print('data[]',data['shapes'][j]['label'])
                                    x = data['shapes'][j]['points'][0][0]
                                    y = data['shapes'][j]['points'][0][1]
                                    if check_points_in_rotated_boxes([x, y], tatou):
                                        print('obb_cls_kpt[obb_i][kpt_i]',obb_cls_kpt[obb_i][kpt_i])
                                        x = round(x / width, 6)
                                        y = round(y / heigth, 6)
                                        line[9 + 3 * (kpt_i + 1) - 3] = str(x)
                                        line[9 + 3 * (kpt_i + 1) - 2] = str(y)
                                        line[9 + 3 * (kpt_i + 1) - 1] = '2'

                    print('line', line)
                    str_line = ' '.join(line)
                    print('str_line', str_line)
                    f.write(str_line + "\n")
        f.close()


if __name__ == '__main__':
    obb_cls=['旋转框类别一','旋转框类别二','旋转框类别三']
    #对应某类别旋转框中关键点，比如旋转框类别一中有4个点，类别二中2个点，三中没有点
    obb_cls_kpt=[['1', '2', '3', '4'],[ 'point1', 'point2'],[]]
    # obb_cls_kpt = [[], [], []]
    json_root = r'D:\Dataset\jsons'  # COCO Object Instance 类型的标注
    ana_txt_save_path = r'D:\Dataset\labels'  # 保存的路径
    image_root=r'D:\Dataset\imgs'


    for filename in os.listdir(json_root):
        json_file=json_root+'\\'+filename
        data = json.load(open(json_file, 'r',encoding='utf-8'))

        image_file=image_root+'\\'+filename.split('.json')[0]+'.jpg'
        # if  os.path.exists(image_file)==False:
        image=cv2.imdecode(np.fromfile(image_file, dtype=np.uint8), -1)
        print(image.shape)
        print('image_file',image_file)
        width=image.shape[1]
        heigth=image.shape[0]
        if not os.path.exists(ana_txt_save_path):
            os.makedirs(ana_txt_save_path)

        save_txt(data,ana_txt_save_path,obb_cls_kpt,obb_cls,width,heigth)

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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158

四、运行命令

yolo obb_pose  mode=train model=ultralytics/cfg/models/v8/yolov8n-obb-pose.yaml  data=ultralytics/cfg/datasets/project/kpt/obb_kpt.yaml batch=8 epochs=100 imgsz=640 workers=0 device=0

 yolo task=obb_pose mode=predict model=runs/obb_pose/train/weights/best.pt  source=images/train
1
2
3