[Other] Support batch eval for detection (#587)

* add onnx_ort_runtime demo

* rm in requirements

* support batch eval
This commit is contained in:
WJJ1995
2022-11-14 21:33:33 +08:00
committed by GitHub
parent 8c79183ea4
commit e8b22283c9

View File

@@ -23,7 +23,8 @@ def eval_detection(model,
ann_file, ann_file,
conf_threshold=None, conf_threshold=None,
nms_iou_threshold=None, nms_iou_threshold=None,
plot=False): plot=False,
batch_size=1):
from .utils import CocoDetection from .utils import CocoDetection
from .utils import COCOMetric from .utils import COCOMetric
import cv2 import cv2
@@ -54,6 +55,8 @@ def eval_detection(model,
start_time = 0 start_time = 0
end_time = 0 end_time = 0
average_inference_time = 0 average_inference_time = 0
im_list = list()
im_id_list = list()
for image_info, i in zip(all_image_info, for image_info, i in zip(all_image_info,
trange( trange(
image_num, desc="Inference Progress")): image_num, desc="Inference Progress")):
@@ -61,19 +64,43 @@ def eval_detection(model,
start_time = time.time() start_time = time.time()
im = cv2.imread(image_info["image"]) im = cv2.imread(image_info["image"])
im_id = image_info["im_id"] im_id = image_info["im_id"]
if batch_size == 1:
if conf_threshold is None and nms_iou_threshold is None: if conf_threshold is None and nms_iou_threshold is None:
result = model.predict(im.copy()) result = model.predict(im.copy())
else: else:
result = model.predict(im, conf_threshold, nms_iou_threshold) result = model.predict(im, conf_threshold, nms_iou_threshold)
pred = { pred = {
'bbox': 'bbox': [[c] + [s] + b
[[c] + [s] + b for b, s, c in zip(result.boxes, result.scores,
for b, s, c in zip(result.boxes, result.scores, result.label_ids) result.label_ids)],
],
'bbox_num': len(result.boxes), 'bbox_num': len(result.boxes),
'im_id': im_id 'im_id': im_id
} }
eval_metric.update(im_id, pred) eval_metric.update(im_id, pred)
else:
im_list.append(im)
im_id_list.append(im_id)
# If the batch_size is not satisfied, the remaining pictures are formed into a batch
if (i + 1) % batch_size != 0 and i != image_num - 1:
continue
if conf_threshold is None and nms_iou_threshold is None:
results = model.batch_predict(im_list)
else:
model.postprocessor.conf_threshold = conf_threshold
model.postprocessor.nms_threshold = nms_iou_threshold
results = model.batch_predict(im_list)
for k in range(len(im_list)):
pred = {
'bbox': [[c] + [s] + b
for b, s, c in zip(results[k].boxes, results[
k].scores, results[k].label_ids)],
'bbox_num': len(results[k].boxes),
'im_id': im_id_list[k]
}
eval_metric.update(im_id_list[k], pred)
im_list.clear()
im_id_list.clear()
if i == image_num - 1: if i == image_num - 1:
end_time = time.time() end_time = time.time()
average_inference_time = round( average_inference_time = round(