【深度学习】Yolov5输出的Yolo格式txt文件转换并写入一个json文件

参考： python将文件内容转为json格式
侵权联系立删

---

前言

将Yolov5预测输出的yolo格式的txt文件转化成json格式，储存在同一个json文件内

---

一、txt和json文件内容说明

txt从左到右为：目标类别，x中心坐标，y中心坐标，目标宽，目标高，置信度
json需要如以下格式存储：

分别为：文件名、标签、预测框坐标（左上角x，左上角y，右下角x，右下角y）、置信度

二、代码

import os

import json
from PIL import Image


# 定义读取txt文件函数
def readFile(filename):
    cata = []
    xm = []
    ym = []
    w = []
    h = []
    score = []
    with open(filename, encoding='utf-8') as lines:
        for line in lines:
            cata.append(line.split(" ")[0])
            xm.append(line.split()[1])
            ym.append(line.split(" ")[2])
            w.append(line.split(" ")[3])
            h.append(line.split(" ")[4])
            score.append(line.split(" ")[5].strip())
    return cata, xm, ym, w, h, score

for j in range(1, 1001):

	# 输出的yolo格式文件
    filename = "/home/lmj/Desktop/zy/runs/detect/exp/labels/test%d.txt" % j
    # 对应的图片路径
    imgpath = "/home/lmj/Desktop/zy/zy_data/test/test%d.jpg" % j
    # 对应图片名
    imgname = 'test%d.jpg' % j

    if os.path.exists('/home/lmj/Desktop/zy/runs/detect/exp/labels/test%d.txt' % j):

        img = Image.open(imgpath)
        imgSize = img.size
        dw = imgSize[0]
        dh = imgSize[1]

        d_1, d_2, d_3, d_4, d_5, d_6 = readFile(filename)


        # 将每个结果以字典的形式存储,再将每个字典存储到列表中
        result = []
        result.clear()
        # 遍历当前的txt文件
        for i in range(len(d_1)):
            # 标签对应的目标（0是fire，1是smoke）
            if int(d_1[i]):
                d_7 = 'smoke'
            else:
                d_7 = 'fire'
            dic = {'name': imgname, 'catagory': d_7, 'bbox': [abs(round(((float(d_2[i]))*dw-(0.5*(float(d_4[i]))*dw)), 2)), abs(round(((float(d_3[i]))*dh-(0.5*(float(d_5[i]))*dh)), 2)), abs(round(((float(d_2[i]))*dw+(0.5*(float(d_4[i]))*dw)), 2)), abs(round(((float(d_3[i]))*dh+(0.5*(float(d_5[i]))*dh)), 2))], 'score': (d_6[i])}
            result.append(dic)

        # 将内容写入到文件中
        newFile = 'Submission.json'
        with open(newFile, 'a') as name:
                for i in range(len(result)):
                    jsonData = json.dumps(result[i], indent=4)
                    # 写入文件
                    name.write(jsonData+','+'\n')

    else:

        abc = []
        dic = {'name': imgname, 'catagory': None, 'bbox': None, 'score': None}
        abc.append(dic)

        with open(newFile, 'a') as name:
                for i in range(len(abc)):
                    jsonData = json.dumps(abc[i], indent=4)
                    # 写入文件
                    name.write(jsonData+','+'\n')