最近很多朋友留言问我如何将检测结果写入json文件并且用于COCO API的评估,之前对于检测结果的格式已经做了简单的说明,这里提供一些简单的函数,直接调用将结果写入即可。
用于COCO API测试的文件格式
HUST小菜鸡:用于COCO API测试的结果的文件格式zhuanlan.zhihu.com使用COCO API进行结果评估
HUST小菜鸡:使用COCO API评估模型在COCO数据集上的结果zhuanlan.zhihu.comCOCO utils给出了一些转换的函数
def det2json(dataset, results):json_results = []for idx in range(len(dataset)):img_id = dataset.img_ids[idx]result = results[idx]for label in range(len(result)):bboxes = result[label]for i in range(bboxes.shape[0]):data = dict()data['image_id'] = img_iddata['bbox'] = xyxy2xywh(bboxes[i])data['score'] = float(bboxes[i][4])data['category_id'] = dataset.cat_ids[label]json_results.append(data)return json_results
def results2json(dataset, results, out_file):result_files = dict()if isinstance(results[0], list):json_results = det2json(dataset, results)result_files['bbox'] = '{}.{}.json'.format(out_file, 'bbox')result_files['proposal'] = '{}.{}.json'.format(out_file, 'bbox')mmcv.dump(json_results, result_files['bbox'])elif isinstance(results[0], tuple):json_results = segm2json(dataset, results)result_files['bbox'] = '{}.{}.json'.format(out_file, 'bbox')result_files['proposal'] = '{}.{}.json'.format(out_file, 'bbox')result_files['segm'] = '{}.{}.json'.format(out_file, 'segm')mmcv.dump(json_results[0], result_files['bbox'])mmcv.dump(json_results[1], result_files['segm'])elif isinstance(results[0], np.ndarray):json_results = proposal2json(dataset, results)result_files['proposal'] = '{}.{}.json'.format(out_file, 'proposal')mmcv.dump(json_results, result_files['proposal'])else:raise TypeError('invalid type of results')return result_files
其他的实现方式也差不多和这个相同
......
省略部分未模型初始化等操作部分if not distributed:model = MMDataParallel(model, device_ids=[0])outputs = single_gpu_test(model, data_loader, args.show, args.save_img, args.save_img_dir)else:model = MMDistributedDataParallel(model.cuda())outputs = multi_gpu_test(model, data_loader, args.tmpdir)res = []for id, boxes in enumerate(outputs):boxes=boxes[0]if type(boxes) == list:boxes = boxes[0]boxes[:, [2, 3]] -= boxes[:, [0, 1]]if len(boxes) > 0:for box in boxes:temp = dict()temp['image_id'] = id+1temp['category_id'] = 1temp['bbox'] = box[:4].tolist()temp['score'] = float(box[4])res.append(temp)with open(args.out, 'w') as f:json.dump(res, f)