yolo txt格式转coco json格式
**问题背景:**下载coco128数据集,使用yolov5模型进行推理并使用pycocotools.cocoeval 对预测结果进行精度计算。
coco128 下载地址:https://tianchi.aliyun.com/dataset/108650
解压缩cocozip之后可以看到如下的目录层级 :
在两个training017目录下,分别包含一个图片jpg文件和label txt文件。
├── images
│ └── train2017
└── 000000000009.jpg
├── labels
│ └── train2017
└── 000000000009.txt
转换代码:
import os
import json
import cv2
import random
import time
from PIL import ImageCOCO_REVERSE_LABEL_MAP = { 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7, 8: 8,9: 9, 10: 10, 11: 11, 12: 13, 13: 14, 14: 15, 15: 16, 16: 17,17: 18, 18: 19, 19: 20, 20: 21, 21: 22, 22: 23, 23: 24, 24: 25,25: 27, 26: 28, 27: 31, 28: 32, 29: 33, 30: 34, 31: 35, 32: 36,33: 37, 34: 38, 35: 39, 36: 40, 37: 41, 38: 42, 39: 43, 40: 44,41: 46, 42: 47, 43: 48, 44: 49, 45: 50, 46: 51, 47: 52, 48: 53,49: 54, 50: 55, 51: 56, 52: 57, 53: 58, 54: 59, 55: 60, 56: 61,57: 62, 58: 63, 59: 64, 60: 65, 61: 67, 62: 70, 63: 72, 64: 73,65: 74, 66: 75, 67: 76, 68: 77, 69: 78, 70: 79, 71: 80, 72: 81,73: 82, 74: 84, 75: 85, 76: 86, 77: 87, 78: 88, 79: 89, 80: 90}id_to_cate = {1: 'person', 2: 'bicycle', 3: 'car', 4: 'motorcycle', 5: 'airplane', 6: 'bus', 7: 'train', 8: 'truck',9: 'boat', 10: 'traffic light', 11: 'fire hydrant', 13: 'stop sign', 14: 'parking meter', 15: 'bench',16: 'bird', 17: 'cat', 18: 'dog', 19: 'horse', 20: 'sheep', 21: 'cow', 22: 'elephant', 23: 'bear',24: 'zebra', 25: 'giraffe', 27: 'backpack', 28: 'umbrella', 31: 'handbag', 32: 'tie', 33: 'suitcase',34: 'frisbee', 35: 'skis', 36: 'snowboard', 37: 'sports ball', 38: 'kite', 39: 'baseball bat',40: 'baseball glove', 41: 'skateboard', 42: 'surfboard', 43: 'tennis racket', 44: 'bottle', 46: 'wine glass',47: 'cup', 48: 'fork', 49: 'knife', 50: 'spoon', 51: 'bowl', 52: 'banana', 53: 'apple', 54: 'sandwich',55: 'orange', 56: 'broccoli', 57: 'carrot', 58: 'hot dog', 59: 'pizza', 60: 'donut', 61: 'cake',62: 'chair', 63: 'couch', 64: 'potted plant', 65: 'bed', 67: 'dining table', 70: 'toilet', 72: 'tv',73: 'laptop', 74: 'mouse', 75: 'remote', 76: 'keyboard', 77: 'cell phone', 78: 'microwave', 79: 'oven',80: 'toaster', 81: 'sink', 82: 'refrigerator', 84: 'book', 85: 'clock', 86: 'vase', 87: 'scissors',88: 'teddy bear', 89: 'hair drier', 90: 'toothbrush'
}def yolotxt2cocojson(save_dir,label_dir,img_dir):categories=[]for id,label in id_to_cate.items():categories.append({'id':id,'name':label,'supercategory':'None'})write_json_context=dict() #写入.json文件的大字典write_json_context['info']= {'description': '', 'url': '', 'version': '', 'year': 2024, 'contributor': 'coco128', 'date_created': '2024-09-06'}write_json_context['licenses']=[{'id':1,'name':None,'url':None}]write_json_context['categories']=categorieswrite_json_context['images']=[]write_json_context['annotations']=[]#接下来的代码主要添加'images'和'annotations'的key值img_list_ori=os.listdir(img_dir) #遍历该文件夹下的所有文件,并将所有文件名添加到列表中label_list_ori=os.listdir(label_dir)# label与image中对应的数据(原始下载coco128数据 image和label并非一一对应)img_list_ori_names = [file.replace('.jpg', '') for file in img_list_ori ]print( "img_list_ori_names = ", len(img_list_ori_names) )label_list_ori_names = [file.replace('.txt', '') for file in label_list_ori ]intersection = [ item for item in img_list_ori_names if item in label_list_ori_names ]print( len(intersection) )img_list = [name+".jpg" for name in intersection ]print( "img_list = ",len(img_list) )for i,img_name in enumerate(img_list):img_path = os.path.join(img_dir,img_name) image = Image.open(img_path) #读取图片,然后获取图片的宽和高W, H = image.sizeimg_context={} #使用一个字典存储该图片信息img_context['file_name']=img_nameimg_context['height']=Himg_context['width']=Wimg_context['date_captured']='2022-07-8'img_context['id']=i #该图片的idimg_context['license']=1img_context['color_url']=''img_context['flickr_url']=''write_json_context['images'].append(img_context) #将该图片信息添加到'image'列表中label_name = img_name.split('.')[0] +'.txt' #获取该图片对应的txt文件with open(os.path.join(label_dir,label_name),'r') as fr:lines=fr.readlines() #读取txt文件的每一行数据,lines2是一个列表,包含了一个图片的所有标注信息for j,line in enumerate(lines):bbox_dict = {} #将每一个bounding box信息存储在该字典中class_id,x,y,w,h=line.strip().split(' ') #获取每一个标注框的详细信息class_id,x, y, w, h = int(class_id), float(x), float(y), float(w), float(h) #将字符串类型转为可计算的int和float类型xmin=(x-w/2)*W #坐标转换ymin=(y-h/2)*Hxmax=(x+w/2)*Wymax=(y+h/2)*Hw=w*Wh=h*Hbbox_dict['id']=i*10000+j #bounding box的坐标信息bbox_dict['image_id']=ibbox_dict['category_id']=COCO_REVERSE_LABEL_MAP[class_id+1] #注意目标类别要加一bbox_dict['iscrowd']=0height,width=abs(ymax-ymin),abs(xmax-xmin)bbox_dict['area']=height*widthbbox_dict['bbox']=[xmin,ymin,w,h]bbox_dict['segmentation']=[]write_json_context['annotations'].append(bbox_dict) #将每一个由字典存储的bounding box信息添加到'annotations'列表中save_name = img_dir.split('/')[-1]save_name = "test0906"name = os.path.join(save_dir,'{}.json'.format(save_name))with open(name,'w') as fw: #将字典信息写入.json文件中json.dump(write_json_context,fw,indent=2,ensure_ascii=False) #加后缀ensure_ascii=False放置写入后中文乱码if __name__ =="__main__":save_dir = "./"label_dir = "./labels/train2017/"img_dir = "./images/train2017/"yolotxt2cocojson(save_dir,label_dir,img_dir)
参考链接:https://blog.csdn.net/a1004550653/article/details/131301909