python 深度学习合并推理出来的txt标签，并生成json

天智杯竞赛格式要求

要求以一个json文件，存储所有结果。

目前我们推理完是生成很多txt文档，而且是拆图之后的，需要将txt中的标签信息转换为大图的，再合成一个json文件，代码如下：

import os
import json
import numpy as np

'''
这个程序是用来合infer推理后生成的txt，因此标签有置信度的分数，同时类别标签在开头
'''

input_dir_path = '/home/rtx2080ti/GeTu/AllTxt/Result'             # 原来存放mask.txt文件的目录；
# os.mkdir('/home/rtx2080ti/GeTu/train/new_mask')                 # 没有目录的话，创建修改后存放的txt目录；
output_dir_path = '/home/rtx2080ti/GeTu/AllTxt/Merge'             # 修改后存放的txt目录；
output_json_path ='/home/rtx2080ti/GeTu/AllTxt/json'

def MergeTxt(input_dir_path):
    for input_filename in os.listdir(input_dir_path):
        print(input_filename)
        # 去除文件名后缀
        name = input_filename[:-4]
        #  得到原图像(也即txt)索引 + 切割高 + 切割宽
        name_list = name.split('_')
        # 得到输出的文件名
        output_filename = name_list[0] + '.txt'    # 原来是一个大图的写在一起
        h = float(name_list[1])                    # 相当于y0
        w = float(name_list[2])                    # 相当于x0
        print(output_filename)
        print(w, h)
        with open(input_dir_path+'/'+input_filename, 'r') as inputfile:
            with open(output_dir_path+'/'+output_filename, 'a') as outputfile:
                for line in inputfile:
                    list1 = line.rstrip('\n').split(' ')
                    print(list1)         # list1是推理出来的小图的标签，分数，和坐标；
                    label = list1[0]     # 这里的label指的是类别标签，在舰船数据集中就为bigship；
                    score = list1[1]     # score为置信度得分；
                    x1 = float(list1[2])
                    y1 = float(list1[3])
                    x2 = float(list1[4])
                    y2 = float(list1[5])
                    x3 = float(list1[6])
                    y3 = float(list1[7])
                    x4 = float(list1[8])
                    y4 = float(list1[9])
                    # 开始转换坐标
                    x1 = float(x1 + w)
                    y1 = float(y1 + h)
                    x2 = float(x2 + w)
                    y2 = float(y2 + h)
                    x3 = float(x3 + w)
                    y3 = float(y3 + h)
                    x4 = float(x4 + w)
                    y4 = float(y4 + h)
                    list2 = [label, score, x1, y1, x2, y2, x3, y3, x4, y4]
                    print(list2)           # list2是大图坐标下的的标签，分数，和坐标；
                    # 开始写入txt，第一列是类别，第二列是置信度分数，后面的是坐标；
                    outputfile.write(str(label) + " " + str(score) + " " + str(x1) + " " + str(y1) + " " + str(x2) + " " +
                                     str(y2) + " " + str(x3) + " " + str(y3) + " " + str(x4) + " " + str(y4) + " " + '\n')
            outputfile.close()
        inputfile.close()
    return outputfile

def CreatJson(output_dir_path):
    with open(os.path.join(output_json_path, 'ship_results.json'), 'a') as outputfile:
        result_list = []                               # 定义一个大list存放所有图片的所有标签；
        for output_filename in os.listdir(output_dir_path):  # 读取输入文件名，写作output是因为是MergeTxt的输出；
            print(output_filename)
            image_name = output_filename[:-4]          # 去掉txt的后缀；
            image_name = image_name + '.png'           # 确定标注文件所归属的图片名称；
            print(image_name)
            with open(output_dir_path + '/' + output_filename, 'r') as inputfile:     # 打开txt文件作为函数的输入文件；
                json_image_name = image_name                 # json中要读到图片名称；
                res_image = {'image_name': json_image_name,  # 定义res_image存放一个图中的所有标签；
                             'labels': []}
                result_list.append(res_image)              # 将res_image的内容，加入到result_list这个表格中,以后都是在res_image后面append；
                for line in inputfile:                     # 遍历一个txt中的每一行；
                    list3 = line.rstrip('\n').split(' ')   # 以空格为间隔读取txt中的一行标注信息；
                    print(list3)
                    json_category_id = list3[0]            # 类别标签，读取txt中目标为哪类；
                    json_confidence = list3[1]             # json_confidence为置信度得分,即为score；
                    json_x1 = float(list3[2])              # 合并后的坐标x1，其他坐标值同理；
                    json_y1 = float(list3[3])
                    json_x2 = float(list3[4])
                    json_y2 = float(list3[5])
                    json_x3 = float(list3[6])
                    json_y3 = float(list3[7])
                    json_x4 = float(list3[8])
                    json_y4 = float(list3[9])
                    # 形成一个只含bounding box的列表；
                    obj_bbox_list = [json_x1, json_y1, json_x2, json_y2, json_x3, json_y3, json_x4, json_y4]
                    obj_dict = {'category_id': json_category_id,
                                'points': np.array(obj_bbox_list).reshape(4, 2).tolist(),  # 将bbox重新reshape为四行两列；
                                'confidence': json_confidence}   # 读置信度；
                    res_image['labels'].append(obj_dict)  # 将obj——dict表示的标签信息，加入到label中，以后都是在labels后面append；
        a = json.dumps(result_list, indent=True)          # 定义一个将result_list生成json结果的命令；
        outputfile.write(a)                               # 将result_list写入到json中；
        inputfile.close()                                 # 关闭输入文件；
    outputfile.close()                                    # 关闭输出文件；
    return outputfile

if __name__ == '__main__':
    MergeTxt(input_dir_path)
    CreatJson(output_dir_path)
    print('已生成ship_results.json文件!')