bop数据合并到COCO
- JSON转TXT
- 重命名
- txt文件中类别信息的转换
JSON转TXT
import json
import os,globcategories = [{"id": 12,"name": "OREO","supercategory": "icbin"},{"id": 16,"name": "Paper Cup","supercategory": "icbin"},{"id": 4,"name": "School Glue","supercategory": "icbin"},{"id": 7,"name": "Straw Cups","supercategory": "icbin"},{"id": 9,"name": "Highland","supercategory": "icbin"},{"id": 10,"name": "Soueakair","supercategory": "icbin"},{"id": 2,"name": "Cheez-it","supercategory": "icbin"},{"id": 1,"name": "Copper Plus","supercategory": "icbin"},{"id": 8,"name": "Stir Stick","supercategory": "icbin"},{"id": 14,"name": "Stanley","supercategory": "icbin"},{"id": 3,"name": "Crayola","supercategory": "icbin"},{"id": 13,"name": "Mirado","supercategory": "icbin"},{"id": 11,"name": "Munchkin","supercategory": "icbin"},{"id": 6,"name": "Greenies","supercategory": "icbin"},{"id": 5,"name": "White Board Cake","supercategory": "icbin"},{"id": 15,"name": "Main Arm","supercategory": "icbin"}]def convert(size, box):dw = 1. / (size[0])dh = 1. / (size[1])x = box[0] + box[2] / 2.0y = box[1] + box[3] / 2.0w = box[2]h = box[3]x = x * dww = w * dwy = y * dhh = h * dhreturn (x, y, w, h)def to_yolo(data_path):json_path=data_path+'/scene_gt_coco.json' save_path = data_path+ '/labels/'json_file = json_path # COCO Object Instance 类型的标注ana_txt_save_path = 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不连续!重新映射一下再输出!print(data['categories'])# # categories = sorted(data['categories'], key=lambda x: x['id'])for i, category in enumerate(categories): # id_map[category['id']] = int(category['id'])id_map[category['id']] = i# 通过事先建表来降低时间复杂度max_id = 0for img in data['images']:max_id = max(max_id, img['id'])# 注意这里不能写作 [[]]*(max_id+1),否则列表内的空列表共享地址img_ann_dict = [[] for i in range(max_id+1)] for i, ann in enumerate(data['annotations']):img_ann_dict[ann['image_id']].append(i)for img in 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.split('/')[-1] + ".txt" # 对应的txt名字,与jpg一致f_txt = open(os.path.join(ana_txt_save_path, ana_txt_name), '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]))'''# 这里可以直接查表而无需重复遍历for ann_id in img_ann_dict[img_id]:ann = data['annotations'][ann_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()print(f'==> coco to yolo images:{len(data["images"])}, save path: {save_path}')def train_val_test(data_path):sets = ['train','val','test']#生成txt的文件名称image_ids = glob.glob(os.path.join(data_path, 'images', '*.jpg'))train_ratio = 0.7 # 训练集比例val_ratio = 0.2 # 验证集比例test_ratio = 0.1 # 测试集比例train_size = int(len(image_ids) * train_ratio)val_size = int(len(image_ids) * val_ratio)test_size = len(image_ids) - train_size - val_sizedata = [image_ids[:train_size], image_ids[train_size:train_size + val_size], image_ids[train_size + val_size:]]for i, image_set in enumerate(sets):image_ids = data[i]list_file = open(data_path+'/%s.txt' % (image_set), 'w')for image_id in image_ids:image_id = image_id.replace('/rgb','/images')list_file.write(image_id + "\n")# convert_annotation(image_id)# 关闭文件list_file.close()print(f'==> train image: {train_size}')print(f'==> valid image: {val_size}')print(f'==> test image: {test_size}')if __name__ == '__main__':data_path = 'H:/Dataset/COCO/train_pbr/000002'to_yolo(data_path)train_val_test(data_path)# print([cat['name'] for cat in categories])
重命名
以00000061*开头
txt文件中类别信息的转换
加79(从0开始,80类的COCO)
import codecs
import ospath = 'H:/Dataset/COCO/train_pbr/000002/labelNew/' # 标签文件train路径
m = os.listdir(path)
# 读取路径下的txt文件
for n in range(0, len(m)):t = codecs.open('H:/Dataset/COCO/train_pbr/000002/labelNew/' + m[n], mode='r', encoding='utf-8')line = t.readline() # 以行的形式进行读取文件list1 = []while line:a = line.split()list1.append(a)line = t.readline()t.close()lt = open('H:/Dataset/COCO/train_pbr/000002/labelNew/' + m[n], "w")for num in range(0, len(list1)):list1[num][0] = str(int(list1[num][0])+79) # 第一列为0时,将0改为1lt.writelines(' '.join(list1[num]) + '\n') # 每个元素以空格间隔,一行元素写完并换行lt.close()print(m[n] + " 修改完成")