3.数据处理-把labelme标注的关键点转成coco格式

# *_* : coding: utf-8 *_*
#####################
#cao
'''
datasets process for object detection project.
for convert customer dataset format to coco data format,
'''

import traceback
import argparse
import datetime
import json
import cv2
import os

__CLASS__ = ['__background__', 'lpr']   # class dictionary, background must be in first index.

def argparser():
    parser = argparse.ArgumentParser("define argument parser for pycococreator!")
    #图片和标签的路径
    parser.add_argument("-r", "--project_path", default="/mnt//", help="path of root directory")
    parser.add_argument("-l", "--label_path", default="/mn/labels",
                        help="path of root directory")
    #图片所在的文件夹，可以支持多个文件夹
    parser.add_argument("-p", "--phase_folder", default=["ceme"], help="datasets path of [train, val, test]")

    #关键点个数
    parser.add_argument("-joint", "--pointNum", default=33,help="point numbers")
    # 是否开启关键点读取，
    parser.add_argument("-po", "--have_points", default=True, help="if have points we will deal it!")
    # 是否开启把点画在图片上验证读取，
    parser.add_argument("-check", "--have_check", default=True, help="if you want to check all points!")

    return parser.parse_args()

def MainProcessing(args):
    '''main process source code.'''
    annotations = {}    # annotations dictionary, which will dump to json format file.
    project_path = args.project_path
    phase_folder = args.phase_folder
    # coco annotations info.
    annotations["info"] = {
        "description": "customer dataset format convert to COCO format",
        "url": "http://cocodataset.org",
        "version": "1.0",
        "year": 2020,
        "contributor": "andy.wei",
        "date_created": "2020/08/24"
    }
    # coco annotations licenses.
    annotations["licenses"] = [{
        "url": "https://www.apache.org/licenses/LICENSE-2.0.html",
        "id": 1,
        "name": "Apache License 2.0"
    }]
    # coco annotations categories.
    annotations["categories"] = []
    for cls, clsname in enumerate(__CLASS__):
        if clsname == '__background__':
            continue
        annotations["categories"].append(
            {
                "supercategory": "object",
                "id": cls,
                "name": clsname
            }
        )
        #保存自己的标签
        for catdict in annotations["categories"]:
            if "lpr" == catdict["name"] and args.have_points:
                catdict["keypoints"] = ["LT1", "LT2", "LT3", "LT4", "LT5","LB1", "LB2", "LB3", "LB4", "LB5",
            "RT1", "RT2", "RT3", "RT4", "RT5","RB1", "RB2", "RB3", "RB4", "RB5","MidLT1","MidLB1","MidRT1","MidRB1",
                                        "1","2","3","4","5","6","7","8","9"]
                catdict["skeleton"] = [[]]

    for phase in phase_folder:
        annotations["images"] = []
        annotations["annotations"] = []
        #label文件夹下txt
        label_path = args.label_path
        #图片的绝对路径
        images_folder = os.path.join(project_path, phase)

        if os.path.isdir(label_path) and os.path.exists(images_folder):
            print("convert datasets {} to coco format!".format(phase))
            step = 0
            for id, line in enumerate(os.listdir(images_folder)):
                print(line)
                #标签名字
                label_name = line.split('.')[0] + '.txt'
                labeltxt = os.path.join(label_path, label_name)
                fd = open(labeltxt, "r")
                if fd:
                    fds = fd.readlines()
                    print(fds)
                    label_info = fds[0].split()
                    print(label_info)

                    #图片的名字
                    image_name = line

                    if fds:

                        bbox = [label_info[1],label_info[2], label_info[3], label_info[4]]
                        cls = int(label_info[0])
                        x1 = float(bbox[0])
                        y1 = float(bbox[1])
                        bw = float(bbox[2]) - float(bbox[0])
                        bh = float(bbox[3]) - float(bbox[1])
                    filename = os.path.join(images_folder, image_name)

                    print(filename)
                    #读取图片名字

                    img = cv2.imread(filename)
                    print(img)
                    if args.have_check:
                        for i in range(args.pointNum):
                            cv2.circle(img,(int(float(label_info[5 + i*3])),int(float(label_info[6 + 3*i]))),2,(0,0,255),2)
                        pthsave ="/mimages/" + str(id) +".jpg"
                        cv2.imwrite(pthsave,img)
                    if img is None:
                        continue
                    height, width, _ = img.shape

                    annotations["images"].append(
                        {
                            "license": 1,
                            "file_name": image_name,
                            "coco_url": "",
                            "height": height,
                            "width": width,
                            "date_captured": datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                            "flickr_url": "",
                            "id": id
                        }
                    )
                    # coco annotations annotations.
                    annotations["annotations"].append(
                        {
                            "id": id,
                            "image_id": id,
                            "category_id": cls,
                            "segmentation": [[]],
                            "area": bw*bh,
                            "bbox": [x1, y1, bw, bh],
                            "iscrowd": 0,
                        }
                    )
                    #是否开启写入关键点
                    if args.have_points:

                        catdict = annotations["annotations"][id]

                            #points = [int(p) for p in label_info[2].split(",")]
                        catdict["keypoints"] = [int(float(label_info[i])) for i in range(5,5+args.pointNum*3)]

                        catdict["num_keypoints"] = args.pointNum



                    step += 1
                    if step % 100 == 0:
                        print("processing {} ...".format(step))
            fd.close()

        else:
            print("WARNNING: file path incomplete, please check!")

        json_path = os.path.join(project_path, "train.json")
        with open(json_path, "w") as f:
            json.dump(annotations, f)


if __name__ == "__main__":
    print("begining to convert customer format to coco format!")
    args = argparser()
    try:
        MainProcessing(args)
    except Exception as e:
        traceback.print_exc()
    print("successful to convert customer format to coco format")