目标检测脚本之mmpose json转yolo txt格式

目标检测脚本之mmpose json转yolo txt格式

一、需求分析

在使用yolopose及yolov8-pose 网络进行人体姿态检测任务时，有时需要标注一些特定场景的中的人型目标数据，用来扩充训练集，提升自己训练模型的效果。因为单纯的人工标注耗时费力，所以可以使用一些开源的大模型如（mmpose）来标注图片。以mmpose为例，对下面图片进行预测后生成的结果图如下所示：

1.预测图片

从预测图片可以看出，预测的结果还是不错的，目标框和关键点的位置预测的还是比较准确的。

2. json 文件

从json 文件中可以看出，包含4个字段：Keyponts、keypoints_score、bbox、bbox_score，注意其中keypoints只包含关键点的(x,y）坐标值，并没有yolo格式中的v值。keypoints_score表示每个关键点的得分，bbox表示目标框的左上角和右下角坐标，bbox_score表示目标框的得分。

二、需求实现

1. 预计结果

如下图所示，希望通过脚本文件可以批量将mmpose生成的json文件转为yolo的txt标签格式

2. 源码实现：

# 处理 mmpose 推理后的json文件，生成coco格式关键点的标签文件
import json
import os
import cv2# 目标检测框 x1y1x2y2 转 为 cls,x_center,y_center,w,h
def convert_xywh(box,image_width,image_height):x1,y1 = box[0],box[1]x2,y2 = box[2],box[3]x = (x2 + x1) /(2*image_width)y = (y2 + y1) /(2*image_height)width = (x2-x1) / image_widthheight = (y2-y1) / image_heightclass_id = 0return [class_id,round(x,4),round(y,4),round(width,4),round(height,4)]#  判断目标的尺寸是否太小，这里设置了5以下返回True，可以自行根据实际情况更改
def is_too_small(label,image_width,image_height):if label[4] * image_width < 5 or label[3] * image_height< 5:return Truedef json2txt(json_dir,image_dir,save_label_dir):"""根据输入的json文件夹，图像文件夹，和保存标签的文件夹，生成coco格式的标签文件。Args:json_dir (str): 存放json文件的文件夹路径。image_dir (str): 存放图像文件的文件夹路径。save_label_dir (str): 保存标签文件的文件夹路径。Returns:None"""print(save_label_dir)if not os.path.exists(save_label_dir):os.makedirs(save_label_dir)json_list = os.listdir(json_dir)for json_file in json_list:image_file = os.path.join(image_dir,json_file.split('.')[0]+'.jpg')img = cv2.imread(image_file)if img is None:continuewidth,height = img.shape[1],img.shape[0]json_path = os.path.join(json_dir,json_file)label_path = os.path.join(save_label_dir,json_file.split('.')[0]+'.txt')with open(json_path,'r',encoding='utf-8') as f:result = json.load(f)  # 读取json文件kepoints = []kepoints_scores = []boxes = []for item in result:kepoints.append(item.get("keypoints"))kepoints_scores.append(item.get("keypoint_scores"))boxes.append(item.get("bbox"))# print(len(boxes))coco_model_kepoints = []coco_boxes = []# 共有多少个目标，也相当于多少组关键点评分for i in range(len(kepoints_scores)):coco_model_kepoints_temp = []#每组关键点有17个，每个关键点有2个值，需要将每个关键点都转换成coco格式，补充为3个值for j in range(len(kepoints_scores[i])):                   if kepoints_scores[i][j] >= 0.45:x = kepoints[i][j][0]/width if kepoints[i][j][0]/width > 0 else 0y = kepoints[i][j][1]/height if kepoints[i][j][1]/height > 0 else 0v = 2else:x = 0y = 0v = 0if x > 1:x = 1if y > 1:y = 1x = round(x,4)y = round(y,4)coco_model_kepoints_temp.append(x)coco_model_kepoints_temp.append(y)coco_model_kepoints_temp.append(v)if all(v == 0 for v in coco_model_kepoints_temp) is False:coco_boxes.append(convert_xywh(boxes[i][0],width,height))coco_model_kepoints.append(coco_model_kepoints_temp)final_label = []   for k in range(len(coco_model_kepoints)):temp_label = []for item1 in coco_boxes[k]:temp_label.append(item1)for item2 in coco_model_kepoints[k]:temp_label.append(item2)final_label.append(temp_label)with open(label_path,'w',encoding='utf-8') as f1:for label in final_label:if is_too_small(label,width,height):continuefor item in label:f1.write(str(item)+' ')f1.write('\n')if __name__ == '__main__':json_dir = '' # 放置使用mmpose 预测出来的图片的json 文件路径image_dir = '' # 放置用于预测的图片路径save_label_dir = '' # 用来保存最终label文件的路径json2txt(json_dir,image_dir,save_label_dir)