def _postprocess(self, cls_outputs, box_outputs, scales, mode='global'):
"""Postprocess class and box predictions."""
if not mode:
return cls_outputs, box_outputs
if mode == 'global':
return postprocess.postprocess_global(self.config.as_dict(),
cls_outputs, box_outputs,
scales)
if mode == 'per_class':
return postprocess.postprocess_per_class(self.config.as_dict(),
cls_outputs, box_outputs,
scales)
if mode == 'combined':
return postprocess.postprocess_combined(self.config.as_dict(),
cls_outputs, box_outputs,
scales)
if mode == 'tflite':
if scales is not None:
# pre_mode should be None for TFLite.
raise ValueError(
'scales not supported for TFLite post-processing')
return postprocess.postprocess_tflite(self.config.as_dict(),
cls_outputs, box_outputs)
raise ValueError('Unsupported postprocess mode {}'.format(mode))
def det_post_process(params: Dict[Any, Any], cls_outputs: Dict[int, tf.Tensor],
box_outputs: Dict[int, tf.Tensor], scales: List[float]):
"""Post preprocessing the box/class predictions.
Args:
params: a parameter dictionary that includes `min_level`, `max_level`,
`batch_size`, and `num_classes`.
cls_outputs: an OrderDict with keys representing levels and values
representing logits in [batch_size, height, width, num_anchors].
box_outputs: an OrderDict with keys representing levels and values
representing box regression targets in [batch_size, height, width,
num_anchors * 4].
scales: a list of float values indicating image scale.
Returns:
detections_batch: a batch of detection results. Each detection is a tensor
with each row as [image_id, ymin, xmin, ymax, xmax, score, class].
"""
if params.get('combined_nms', None):
# Use combined version for dynamic batch size.
nms_boxes, nms_scores, nms_classes, _ = postprocess.postprocess_combined(
params, cls_outputs, box_outputs, scales)
else:
nms_boxes, nms_scores, nms_classes, _ = postprocess.postprocess_global(
params, cls_outputs, box_outputs, scales)
batch_size = tf.shape(cls_outputs[params['min_level']])[0]
img_ids = tf.expand_dims(
tf.cast(tf.range(0, batch_size), nms_scores.dtype), -1)
detections = [
img_ids * tf.ones_like(nms_scores),
nms_boxes[:, :, 0],
nms_boxes[:, :, 1],
nms_boxes[:, :, 2],
nms_boxes[:, :, 3],
nms_scores,
nms_classes,
]
return tf.stack(detections, axis=-1, name='detections')