当前位置：首页 > news >正文

yolov8旋转框+关键点检测

news 来源：原创 2024/9/20 8:12:50

一、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.0y = (box[1] + box[3] )/ 2.0w = 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 boxesArgs:points (tensor): (1, L, 2) anchor pointsboxes (tensor): [B, N, 5] gt_bboxeseps (float): default 1e-9Returns: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 - aad = 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 productap_dot_ab = np.sum(ap * ab)# [B, N, L] dot productap_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_boxdef 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个kptfor i in range(len(data['shapes'])):len_line = 9 + max(len(sublist) for sublist in obb_cls_kpt) * 3# len_line=9line = ['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=y3tatou = [[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+'\\'+filenamedata = 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)

四、运行命令

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=0yolo task=obb_pose mode=predict model=runs/obb_pose/train/weights/best.pt  source=images/train