将coco中的关键点转为yolo关键点格式

更改txt文件中的每行第一个数字,yolo的label第一个数字是类别,有时候需要更改这个数字

import os

dir_root = r'路径'
file_name = os.listdir(dir_root)

for dir_dir in file_name:
    file_n = dir_root + "\\" + dir_dir
    print(file_n)
    with open(file_n, 'r+') as f:
        lines = f.readlines()
        line_save = ""
        for line in lines:
            line = list(line)
            if line[0] == "5":
                line[0] = "1"
                line = "".join(line)
                print(line)

            elif line[0] == "3":
                line[0] = "2"
                line = "".join(line)

            elif line[0] == "4":
                line[0] = "3"
                line = "".join(line)

            elif line[0] == "0":
                line[0] = "0"
                line = "".join(line)
            line_save = line_save + line
        print(line_save)

    with open(file_n, "w") as f1:
        f1.write(line_save)

更改一个文件夹下的所有文件的文件名(全部重命名)

import os

filename = r"文件夹"

x = os.listdir(filename)
m = 1
for file in x:
    file = os.path.join(filename, file)
    os.chdir(filename)
    os.rename(file, "{}.txt".format(m))
    m += 1

下面的是将COCO 格式的数据集转化为 YOLO 格式的数据集,这个代码是从网上找到的,但是它这里没有将关键点也转换出来,所以又写了一个转化关键点的

# 格式转换.py
import json
from tqdm import tqdm
import argparse

parser = argparse.ArgumentParser()
#这里根据自己的json文件位置,换成自己的就行
parser.add_argument('--json_path', default=r'F:\debug-ed\gai\spacecraft\annotations\person_keypoints_train2017.json',type=str, help="input: coco format(json)")
#这里设置.txt文件保存位置
parser.add_argument('--save_path', default=r'F:\debug-ed\gai\spacecraft\cc\baocun', type=str, help="specify where to save the output dir of labels")
arg = parser.parse_args()

def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = box[0] + box[2] / 2.0
    y = box[1] + box[3] / 2.0
    w = box[2]
    h = box[3]
#round函数确定(xmin, ymin, xmax, ymax)的小数位数
    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)

if True:
    json_file = arg.json_path # COCO Object Instance 类型的标注
    ana_txt_save_path = arg.save_path  # 保存的路径

    data = json.load(open(json_file, 'r'))
    if not os.path.exists(ana_txt_save_path):
        os.makedirs(ana_txt_save_path)

    id_map = {} # coco数据集的id不连续!重新映射一下再输出!
    with open(os.path.join(ana_txt_save_path, 'classes.txt'), 'w') as f:
        # 写入classes.txt
        for i, category in enumerate(data['categories']):
            f.write(f"{category['name']}\n")
            id_map[category['id']] = i
    # print(id_map)
    # 这里需要根据自己的需要,更改写入图像相对路径的文件位置。
    list_file = open(os.path.join(ana_txt_save_path, 'val.txt'), 'w')
    
    for img in tqdm(data['images']):
        filename = img["file_name"]
        img_width = img["width"]
        img_height = img["height"]
        img_id = img["id"]
        head, tail = os.path.splitext(filename)
        ana_txt_name = head + ".txt"  # 对应的txt名字,与jpg一致
        f_txt = open(os.path.join(ana_txt_save_path, str(img_id)+".txt"), 'w')
        for ann in data['annotations']:
            if ann['image_id'] == img_id:
                box = convert((img_width, img_height), ann["bbox"])
                f_txt.write("%s %s %s %s %s\n" % (id_map[ann["category_id"]], box[0], box[1], box[2], box[3]))
        f_txt.close()
        # 将图片的相对路径写入train2017或val2017的路径
        list_file.write('val/%s.jpg\n' %(head))
    list_file.close()
# COCO 格式的数据集转化为 YOLO 格式的数据集
# --json_path 输入的json文件路径
# --save_path 保存的文件夹名字,默认为当前目录下的labels。
import numpy as np
import os
import json
from tqdm import tqdm
import argparse
from 转换格式 import *

parser = argparse.ArgumentParser()
# 这里根据自己的json文件位置,换成自己的就行
parser.add_argument('--json_path', default=r'F:\debug-ed\gai\spacecraft\annotations\person_keypoints_train2017.json',
                    type=str,
                    help="input: coco format(json)")
# 这里设置.txt文件保存位置
parser.add_argument('--save_path', default=r'F:\debug-ed\gai\spacecraft\cc\baocun', type=str,
                    help="specify where to save the output dir of labels")
arg = parser.parse_args()

xiaoshu = 10 ** 6  # 保存几位小数, 6位
cun = r'F:\debug-ed\gai\spacecraft\cc\baocun'

if __name__ == '__main__':
    json_file = arg.json_path  # COCO Object Instance 类型的标注
    ana_txt_save_path = arg.save_path  # 保存的路径

    data = json.load(open(json_file, 'r'))
    if not os.path.exists(ana_txt_save_path):
        os.makedirs(ana_txt_save_path)

    id_map = {}  # coco数据集的id不连续!重新映射一下再输出!
    with open(os.path.join(ana_txt_save_path, 'classes.txt'), 'w') as f:
        # 写入classes.txt
        for i, category in enumerate(data['categories']):
            f.write(f"{category['name']}\n")
            id_map[category['id']] = i
    # print(id_map)
    # 这里需要根据自己的需要,更改写入图像相对路径的文件位置。
    list_file = open(os.path.join(ana_txt_save_path, 'val.txt'), 'w')

    for annotations in tqdm(data['annotations']):
        # filename = annotations["file_name"]
        img_width = annotations["width"]
        img_height = annotations["height"]
        img_id = annotations["image_id"]
        keypoints = annotations["keypoints"]

        # print(keypoints)
        arry_x = np.zeros([17, 1])
        num_1 = 0
        for x in keypoints[0:51:3]:
            arry_x[num_1, 0] = int((x / img_width) * xiaoshu) / xiaoshu
            num_1 += 1

        arry_y = np.zeros([17, 1])
        num_2 = 0
        for y in keypoints[1:51:3]:
            arry_y[num_2, 0] = int((y / img_height) * xiaoshu) / xiaoshu
            num_2 += 1

        arry_v = np.zeros([17, 1])
        num_3 = 0
        for v in keypoints[2:51:3]:
            arry_v[num_3, 0] = v
            num_3 += 1

        list_1 = []
        num_4 = 0
        for i in range(17):
            list_1.append(float(arry_x[num_4]))
            list_1.append(float(arry_y[num_4]))
            list_1.append(float(arry_v[num_4]))
            num_4 += 1


        fil = img_id

        cun_1 = os.path.join(cun, str(fil))
        cun_2 = cun_1 + ".txt"

        with open(cun_2, "r") as f:
            read_lines = f.readlines()
        x_1 = read_lines[0].strip().split(",")
        x_2 = x_1 + list_1
        list_xin = list(map(str, x_2))
        list_xin = " ".join(list_xin)
        with open(cun_2, "w") as f:
            f.write(list_xin)

    for img in tqdm(data['images']):
        filename = img["file_name"]
        img_width = img["width"]
        img_height = img["height"]
        img_id = img["id"]
        head, tail = os.path.splitext(filename)
        ana_txt_name = head + ".txt"  # 对应的txt名字,与jpg一致
        os.chdir(cun)
        # os.rename(ana_txt_name, str(img_id) + ".txt")
        os.rename(str(img_id) + ".txt", ana_txt_name)

将第一个代码保存在转换格式.py中,第二个代码要和第一个代码放在一个文件夹下
然后只运行第二个代码就可以了,为了省事就没有精简,这里检测类别只有一个(person),如果多类别的话需要修改,第二个代码中的[1:51:3],这里51是每张图片中关键点keypoints中包含的数据长度是51,arry_v = np.zeros([17, 1]),这里的17就是51/3,有空在对其精简下。

你可能感兴趣的:(python,开发语言)