def handle_prediction(prediction, image_file, image, image_shape, anchors, class_names, model_image_size, elim_grid_sense, v5_decode):
start = time.time()
if len(anchors) == 5:
# YOLOv2 use 5 anchors and have only 1 prediction
assert len(prediction) == 1, 'invalid YOLOv2 prediction number.'
boxes, classes, scores = yolo2_postprocess_np(prediction[0], image_shape, anchors, len(class_names), model_image_size, elim_grid_sense=elim_grid_sense)
else:
if v5_decode:
boxes, classes, scores = yolo5_postprocess_np(prediction, image_shape, anchors, len(class_names), model_image_size, elim_grid_sense=True) #enable "elim_grid_sense" by default
else:
boxes, classes, scores = yolo3_postprocess_np(prediction, image_shape, anchors, len(class_names), model_image_size, elim_grid_sense=elim_grid_sense)
end = time.time()
print("PostProcess time: {:.8f}ms".format((end - start) * 1000))
print('Found {} boxes for {}'.format(len(boxes), image_file))
for box, cls, score in zip(boxes, classes, scores):
xmin, ymin, xmax, ymax = box
print("Class: {}, Score: {}, Box: {},{}".format(class_names[cls], score, (xmin, ymin), (xmax, ymax)))
colors = get_colors(class_names)
image = draw_boxes(image, boxes, classes, scores, class_names, colors)
Image.fromarray(image).show()
return
def predict(self, image_data, image_shape):
num_anchors = len(self.anchors)
if self.model_type.startswith(
'scaled_yolo4_') or self.model_type.startswith('yolo5_'):
# Scaled-YOLOv4 & YOLOv5 entrance, enable "elim_grid_sense" by default
out_boxes, out_classes, out_scores = yolo5_postprocess_np(
self.yolo_model.predict(image_data),
image_shape,
self.anchors,
len(self.class_names),
self.model_image_size,
max_boxes=100,
confidence=self.score,
iou_threshold=self.iou,
elim_grid_sense=True)
elif self.model_type.startswith('yolo3_') or self.model_type.startswith('yolo4_') or \
self.model_type.startswith('tiny_yolo3_') or self.model_type.startswith('tiny_yolo4_'):
# YOLOv3 & v4 entrance
out_boxes, out_classes, out_scores = yolo3_postprocess_np(
self.yolo_model.predict(image_data),
image_shape,
self.anchors,
len(self.class_names),
self.model_image_size,
max_boxes=100,
confidence=self.score,
iou_threshold=self.iou,
elim_grid_sense=self.elim_grid_sense)
elif self.model_type.startswith(
'yolo2_') or self.model_type.startswith('tiny_yolo2_'):
# YOLOv2 entrance
out_boxes, out_classes, out_scores = yolo2_postprocess_np(
self.yolo_model.predict(image_data),
image_shape,
self.anchors,
len(self.class_names),
self.model_image_size,
max_boxes=100,
confidence=self.score,
iou_threshold=self.iou,
elim_grid_sense=self.elim_grid_sense)
else:
raise ValueError('Unsupported model type')
return out_boxes, out_classes, out_scores