目标检测脚本之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. 源码实现：
python
# 处理 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_width
height = (y2-y1) / image_height
class_id = 0
return [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 True
def 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:
continue
width,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 0
y = kepoints[i][j][1]/height if kepoints[i][j][1]/height > 0 else 0
v = 2
else:
x = 0
y = 0
v = 0
if x > 1:
x = 1
if y > 1:
y = 1
x = 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):
continue
for 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)