要求以一个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文件!')
欢迎分享,转载请注明来源:内存溢出
评论列表(0条)