paddleDetection框架如何将模型预测结果保存成json文件
之前我用mmdetection的时候，发现自带的脚本可以把预测结果转为json文件。前面我也用的yolov4代码只用了pytorch框架，根据预测代码修改了一份将模型预测结果保存成json文件的代码。具体参考我的这篇博客将yolov4预测结果保存为json文件。
最近我学了一下paddleDetection框架，里面集成了很多常用的网络模型，用起来很方便，但是我想将模型预测出来的bbox结果保存成json文件，找了一下官方文档，似乎没找到已经写好的脚本，所以我又根据tools/infer.py以及它所引用的代码，改写了一份可以把模型预测结果保存成json文件的代码。
具体思路
具体思路如下：
先从tools/infer.py看起，从run这个函数可以看到调用了ppdet/engine文件夹里面的trainer.py里面Trainer类的predict方法将检测框画在原图上，保存在output文件夹下。有点坑的是，这里的ppdet并不是代码里面的ppdet文件夹，而是安装paddleDetection环境下的ppdet文件夹。所以如果你修改了ppdet文件夹里面的代码运行没有效果记得去找paddle21/lib/python3.8/site-packages/paddledet-2.0.1-py3.8.egg/ppdet/文件夹修改。
from ppdet.engine import Trainer
def run(FLAGS, cfg):
# build trainer
trainer = Trainer(cfg, mode='test')
# load weights
trainer.load_weights(cfg.weights)
# get inference images
images = get_test_images(FLAGS.infer_dir, FLAGS.infer_img)
# inference
trainer.predict(
images,
draw_threshold=FLAGS.draw_threshold,
output_dir=FLAGS.output_dir,
save_txt=FLAGS.save_txt)
接着去看predict函数的具体实现：
from ppdet.utils.visualizer import visualize_results, save_result
def predict(self,
images,
draw_threshold=0.5,
output_dir='output',
save_txt=False):
self.dataset.set_images(images)
loader = create('TestReader')(self.dataset, 0)
imid2path = self.dataset.get_imid2path()
anno_file = self.dataset.get_anno()
clsid2catid, catid2name = get_categories(self.cfg.metric, anno_file)
# Run Infer
self.status['mode'] = 'test'
self.model.eval()
for step_id, data in enumerate(loader):
self.status['step_id'] = step_id
# forward
outs = self.model(data)
for key in ['im_shape', 'scale_factor', 'im_id']:
outs[key] = data[key]
for key, value in outs.items():
outs[key] = value.numpy()
batch_res = get_infer_results(outs, clsid2catid)
bbox_num = outs['bbox_num']
start = 0
for i, im_id in enumerate(outs['im_id']):
image_path = imid2path[int(im_id)]
image = Image.open(image_path).convert('RGB')
self.status['original_image'] = np.array(image.copy())
end = start + bbox_num[i]
bbox_res = batch_res['bbox'][start:end] \
if 'bbox' in batch_res else None
mask_res = batch_res['mask'][start:end] \
if 'mask' in batch_res else None
segm_res = batch_res['segm'][start:end] \
if 'segm' in batch_res else None
image = visualize_results(image, bbox_res, mask_res, segm_res,
int(outs['im_id']), catid2name,
draw_threshold)
self.status['result_image'] = np.array(image.copy())
if self._compose_callback:
self._compose_callback.on_step_end(self.status)
# save image with detection
save_name = self._get_save_image_name(output_dir, image_path)
logger.info("Detection bbox results save in {}".format(
save_name))
image.save(save_name, quality=95)
if save_txt:
save_path = os.path.splitext(save_name)[0] + '.txt'
save_result(save_path, bbox_res, catid2name, draw_threshold)
start = end
我们可以看到，如果将save_txt参数，将会调用ppdet/utils/visualizer.py里面的save_result函数，在output文件夹下生成对应图片的预测信息txt文。
def save_result(save_path, bbox_res, catid2name, threshold):
"""
save result as txt
"""
with open(save_path, 'w') as f:
for dt in bbox_res:
catid, bbox, score = dt['category_id'], dt['bbox'], dt['score']
if score < threshold:
continue
# each bbox result as a line
# for rbox: classname score x1 y1 x2 y2 x3 y3 x4 y4
# for bbox: classname score x1 y1 w h
bbox_pred = '{} {} '.format(catid2name[catid], score) + ' '.join(
[str(e) for e in bbox])
f.write(bbox_pred + '\n')
由此，我们可以添加一个变量save_json，在trainer.py中参考predict函数写一个save_json函数，调用visualizer.py的save_result修改一下，写成save_rusult_json函数。想法是在infer.py里面的run函数那里，从test.json文件把所有测试集的图片信息读进来，然后根据image_name每次取一张，通过调用save_result_json函数和save_json函数将预测框信息保存在content_json列表里，然后调用run函数将content_json列表里面的值写进一个json文件里面。
run函数
def run(FLAGS, cfg):
# build trainer
trainer = Trainer(cfg, mode='test')
# load weights
trainer.load_weights(cfg.weights)
# get inference images
# images = get_test_images(FLAGS.infer_dir, FLAGS.infer_img)
path = "./dataset/coco/test2017/"
path_json = "./dataset/coco/annotations/instances_test2017.json"
save_path = 'result/result.json'
content_json=[]
f = open(path_json, 'r', encoding='utf-8')
s = f.read()
rest = json.loads(s)
#a = -1
for i in rest['images']:
#a+=1
#if a<1501:
# continue
file_name = i["file_name"]
image_id = i["id"]
FLAGS.infer_img=path+file_name
image = get_test_images(FLAGS.infer_dir, FLAGS.infer_img)
content_json.extend(trainer.save_json(
image,
image_id,
draw_threshold=FLAGS.draw_threshold,
output_dir=FLAGS.output_dir,
save_json=FLAGS.save_json))
#if a==1998: # 会在1998后面多出一个],记得删除再提交
# with open(save_path, 'a') as f:
# json.dump(content_json, f)
# break
f.close()
with open(save_path, 'w') as f:
json.dump(content_json, f)
save_json函数
def save_json(self,
image,
image_id,
draw_threshold=0.5,
output_dir='result',
save_json=True):
self.dataset.set_images(image)
loader = create('TestReader')(self.dataset, 0)
imid2path = self.dataset.get_imid2path()
anno_file = self.dataset.get_anno()
clsid2catid, catid2name = get_categories(self.cfg.metric, anno_file)
# Run Infer
self.status['mode'] = 'test'
self.model.eval()
content_json = []
for step_id, data in enumerate(loader):
self.status['step_id'] = step_id
# forward
outs = self.model(data)
for key in ['im_shape', 'scale_factor', 'im_id']:
outs[key] = data[key]
for key, value in outs.items():
outs[key] = value.numpy()
batch_res = get_infer_results(outs, clsid2catid)
bbox_num = outs['bbox_num']
start = 0
for i, im_id in enumerate(outs['im_id']):
image_path = imid2path[int(im_id)]
image = Image.open(image_path).convert('RGB')
self.status['original_image'] = np.array(image.copy())
end = start + bbox_num[i]
bbox_res = batch_res['bbox'][start:end] \
if 'bbox' in batch_res else None
mask_res = batch_res['mask'][start:end] \
if 'mask' in batch_res else None
segm_res = batch_res['segm'][start:end] \
if 'segm' in batch_res else None
if save_json:
content_json.extend(save_result_json(bbox_res, image_id, draw_threshold))
start = end
return content_json
save_result_json函数
def save_result_json(bbox_res, image_id, threshold):
"""
save result as json
"""
content_json=[]
for dt in bbox_res:
catid, bbox, score = dt['category_id'], dt['bbox'], dt['score']
if score < threshold:
continue
# each bbox result as a line
# for rbox: classname score x1 y1 x2 y2 x3 y3 x4 y4
# for bbox: classname score x1 y1 w h
content_dic = {
"image_id":image_id,
"category_id": catid,
"bbox": bbox,
"score": score
}
# print(content_dic)
content_json.append(content_dic)
# with open(save_path, 'w') as f:
# json.dump(content_json, f)
return content_json
这样运行python tools/infer.py -c configs/ppyolo/ppyolov2_r50vd_dcn_365e_coco.yml -o weights=output/ppyolov2_r50vd_dcn_365e_coco/1000.pdparams就可以生成json结果文件了。
若还要在json文件最后加上类别信息，在content_json列表后面加上即可：
#cat={}
#info = [{"id": 1, "name": "rect_eye", "supercategory": "none"}, {"id": 2, "name": "sphere_eye", "supercategory": "none"}, {"id": 3, "name": "box_eye", "supercategory": "none"}]
#cat["categories"]=info
#content_json.extend(cat)
#content_json.extend(cat["categories"])
但是程序运行一会可能会报内存不够的错误，我是直接打开保存的json文件，看看写到那里了，在run函数里面加个计数器，跳过前面的图片信息，由于我们的写入方式是追加模式a，所以在次运行程序可以接着把后面的图片信息补上。内存不够大，又懒得改框架里面具体检测的代码，只能用这种笨方法了。我测1999张图片，跑了三次终于跑完了。如果有更好的实现方式，欢迎交流指正。