目标检测脚本之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_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)