official/vision/evaluation/instance_metrics.py
# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Metrics for instance detection & segmentation."""
from typing import Any, Dict, Optional, Tuple, Union
import numpy as np
import tensorflow as tf, tf_keras
from official.vision.ops import box_ops
from official.vision.ops import mask_ops
class AveragePrecision(tf_keras.layers.Layer):
"""The algorithm which computes average precision from P-R curve."""
def __init__(self, *args, **kwargs):
# Enforce the `AveragePrecision` to operate in `float32` given the
# implementation requirements.
super().__init__(*args, dtype=tf.float32, **kwargs)
def call(self, precisions, recalls):
"""Computes average precision."""
raise NotImplementedError
class COCOAveragePrecision(AveragePrecision):
"""Average precision in COCO style.
In COCO, AP is defined as the mean of interpolated precisions at a set of 101
equally spaced recall points [0, 0.01, ..., 1]. For each recall point r,
the precision is interpolated to the maximum precision with corresponding
recall r' >= r.
The VOC challenges before 2010 used the similar method, but only 11 recall
points [0, 0.1, ..., 1].
"""
def __init__(
self, num_recall_eval_points: int = 101, recalls_desc: bool = False
):
"""Initialization for COCOAveragePrecision.
Args:
num_recall_eval_points: the number of equally spaced recall points used
for interpolating the precisions.
recalls_desc: If true, the recalls are in descending order.
"""
super().__init__()
self._num_recall_eval_points = num_recall_eval_points
self._recalls_desc = recalls_desc
def get_config(self) -> Dict[str, Any]:
return {
'num_recall_eval_points': self._num_recall_eval_points,
'recalls_desc': self._recalls_desc,
}
def call(self, precisions: tf.Tensor, recalls: tf.Tensor) -> tf.Tensor:
"""Computes average precision.
Args:
precisions: a tensor in shape (dim_0, ..., num_confidences) which stores a
list of precision values at different confidence thresholds with
arbitrary numbers of leading dimensions.
recalls: a tensor in shape (dim_0, ..., num_confidences) which stores a
list of recall values at different confidence threshold with arbitrary
numbers of leading dimensions.
Returns:
A tensor in shape (dim_0, ...), which stores the area under P-R curve.
"""
p = precisions
r = recalls
if not isinstance(p, tf.Tensor):
p = tf.convert_to_tensor(p)
if not isinstance(r, tf.Tensor):
r = tf.convert_to_tensor(r)
if self._recalls_desc:
p = tf.reverse(p, axis=[-1])
r = tf.reverse(r, axis=[-1])
r_eval_points = tf.linspace(0.0, 1.0, self._num_recall_eval_points)
# (dim_0, ..., num_recall_eval_points)
# For each recall eval point, the precision is interpolated to the maximum
# precision with corresponding recall >= the recall eval point.
p_max = tf.reduce_max(
p[..., tf.newaxis, :]
* tf.cast(
r[..., tf.newaxis, :] >= r_eval_points[:, tf.newaxis], dtype=p.dtype
),
axis=-1,
)
# (dim_0, ...)
return tf.reduce_mean(p_max, axis=-1)
class VOC2010AveragePrecision(AveragePrecision):
"""Average precision in VOC 2010 style.
Since VOC 2010, first compute an approximation of the measured P-R curve
with precision monotonically decreasing, by setting the precision for recall
r to the maximum precision obtained for any recall r' >= r. Then compute the
AP as the area under this curve by numerical integration.
"""
def __init__(self, recalls_desc: bool = False):
"""Initialization for VOC10AveragePrecision.
Args:
recalls_desc: If true, the recalls are in descending order.
"""
super().__init__()
self._recalls_desc = recalls_desc
def get_config(self) -> Dict[str, Any]:
return {
'recalls_desc': self._recalls_desc,
}
def call(self, precisions: tf.Tensor, recalls: tf.Tensor) -> tf.Tensor:
"""Computes average precision.
Args:
precisions: a tensor in shape (dim_0, ..., num_confidences) which stores a
list of precision values at different confidence thresholds with
arbitrary numbers of leading dimensions.
recalls: a tensor in shape (dim_0, ..., num_confidences) which stores a
list of recall values at different confidence threshold with arbitrary
numbers of leading dimensions.
Returns:
A tensor in shape (dim_0, ...), which stores the area under P-R curve.
"""
p = precisions
r = recalls
if not isinstance(p, tf.Tensor):
p = tf.convert_to_tensor(p)
if not isinstance(r, tf.Tensor):
r = tf.convert_to_tensor(r)
if self._recalls_desc:
p = tf.reverse(p, axis=[-1])
r = tf.reverse(r, axis=[-1])
axis_indices = list(range(len(p.get_shape())))
# Transpose to (num_confidences, ...), because tf.scan only applies to the
# first dimension.
p = tf.transpose(p, np.roll(axis_indices, 1))
# Compute cumulative maximum in reverse order.
# For example, the reverse cumulative maximum of [5,6,3,4,2,1] is
# [6,6,4,4,2,1].
p = tf.scan(
tf.maximum, elems=p, initializer=tf.reduce_min(p, axis=0), reverse=True
)
# Transpose back to (..., num_confidences)
p = tf.transpose(p, np.roll(axis_indices, -1))
# Prepend 0 to r and compute the delta.
r = tf.concat([tf.zeros_like(r[..., 0:1]), r], axis=-1)
delta_r = tf.roll(r, shift=-1, axis=-1) - r
return tf.reduce_sum(p * delta_r[..., :-1], axis=-1)
class MatchingAlgorithm(tf_keras.layers.Layer):
"""The algorithm which matches detections to ground truths."""
def __init__(self, *args, **kwargs):
# Enforce the `MachingAlgorithm` to operate in `float32` given the
# implementation requirements.
super().__init__(*args, dtype=tf.float32, **kwargs)
def call(
self,
detection_to_gt_ious: tf.Tensor,
detection_classes: tf.Tensor,
detection_scores: tf.Tensor,
gt_classes: tf.Tensor,
):
"""Matches detections to ground truths."""
raise NotImplementedError
class COCOMatchingAlgorithm(MatchingAlgorithm):
"""The detection matching algorithm used in COCO."""
def __init__(self, iou_thresholds: Tuple[float, ...]):
"""Initialization for COCOMatchingAlgorithm.
Args:
iou_thresholds: a list of IoU thresholds.
"""
super().__init__()
self._iou_thresholds = iou_thresholds
def get_config(self) -> Dict[str, Any]:
return {
'iou_thresholds': self._iou_thresholds,
}
def call(
self,
detection_to_gt_ious: tf.Tensor,
detection_classes: tf.Tensor,
detection_scores: tf.Tensor,
gt_classes: tf.Tensor,
) -> Tuple[tf.Tensor, tf.Tensor]:
"""Matches detections to ground truths.
This is the matching algorithm used in COCO. First, sort all the detections
based on the scores from high to low. Then for each detection, iterates
through all ground truth. The unmatched ground truth with the highest IoU
greater than the threshold is matched to the detection.
Args:
detection_to_gt_ious: a tensor in shape of (batch_size, num_detections,
num_gts) which stores the IoUs for each pair of detection and ground
truth.
detection_classes: a tensor in shape of (batch_size, num_detections) which
stores the classes of the detections.
detection_scores: a tensor in shape of (batch_size, num_detections) which
stores the scores of the detections.
gt_classes: a tensor in shape of (batch_size, num_gts) which stores the
classes of the ground truth boxes.
Returns:
Two bool tensors in shape of (batch_size, num_detections,
num_iou_thresholds) and (batch_size, num_gts, num_iou_thresholds) which
indicates whether the detections and ground truths are true positives at
different IoU thresholds.
"""
batch_size = tf.shape(detection_classes)[0]
num_detections = detection_classes.get_shape()[1]
num_gts = gt_classes.get_shape()[1]
num_iou_thresholds = len(self._iou_thresholds)
# (batch_size, num_detections)
sorted_detection_indices = tf.argsort(
detection_scores, axis=1, direction='DESCENDING'
)
# (batch_size, num_detections)
sorted_detection_classes = tf.gather(
detection_classes, sorted_detection_indices, batch_dims=1
)
# (batch_size, num_detections, num_gts)
sorted_detection_to_gt_ious = tf.gather(
detection_to_gt_ious, sorted_detection_indices, batch_dims=1
)
init_loop_vars = (
0, # i: the loop counter
tf.zeros(
[batch_size, num_detections, num_iou_thresholds], dtype=tf.bool
), # detection_is_tp
tf.zeros(
[batch_size, num_gts, num_iou_thresholds], dtype=tf.bool
), # gt_is_tp
)
def _match_detection_to_gt_loop_body(
i: int, detection_is_tp: tf.Tensor, gt_is_tp: tf.Tensor
) -> Tuple[int, tf.Tensor, tf.Tensor]:
"""Iterates the sorted detections and matches to the ground truths."""
# (batch_size, num_gts)
gt_ious = sorted_detection_to_gt_ious[:, i, :]
# (batch_size, num_gts, num_iou_thresholds)
gt_matches_detection = (
# Ground truth is not matched yet.
~gt_is_tp
# IoU is greater than the threshold.
& (gt_ious[:, :, tf.newaxis] > self._iou_thresholds)
# Classes are matched.
& (
(sorted_detection_classes[:, i][:, tf.newaxis] == gt_classes)
& (gt_classes > 0)
)[:, :, tf.newaxis]
)
# Finds the matched ground truth with max IoU.
# If there is no matched ground truth, the argmax op will return index 0
# in this step. It's fine because it will be masked out in the next step.
# (batch_size, num_iou_thresholds)
matched_gt_with_max_iou = tf.argmax(
tf.cast(gt_matches_detection, gt_ious.dtype)
* gt_ious[:, :, tf.newaxis],
axis=1,
output_type=tf.int32,
)
# (batch_size, num_gts, num_iou_thresholds)
gt_matches_detection &= tf.one_hot(
matched_gt_with_max_iou,
depth=num_gts,
on_value=True,
off_value=False,
axis=1,
)
# Updates detection_is_tp
# Map index back to the unsorted detections.
# (batch_size, num_detections, num_iou_thresholds)
detection_is_tp |= (
tf.reduce_any(gt_matches_detection, axis=1, keepdims=True)
& tf.one_hot(
sorted_detection_indices[:, i],
depth=num_detections,
on_value=True,
off_value=False,
axis=-1,
)[:, :, tf.newaxis]
)
detection_is_tp.set_shape([None, num_detections, num_iou_thresholds])
# Updates gt_is_tp
# (batch_size, num_gts, num_iou_thresholds)
gt_is_tp |= gt_matches_detection
gt_is_tp.set_shape([None, num_gts, num_iou_thresholds])
# Returns the updated loop vars.
return (i + 1, detection_is_tp, gt_is_tp)
_, detection_is_tp_result, gt_is_tp_result = tf.while_loop(
cond=lambda i, *_: i < num_detections,
body=_match_detection_to_gt_loop_body,
loop_vars=init_loop_vars,
parallel_iterations=32,
maximum_iterations=num_detections,
)
return detection_is_tp_result, gt_is_tp_result
def _shift_and_rescale_boxes(
boxes: tf.Tensor,
output_boundary: Tuple[int, int],
) -> tf.Tensor:
"""Shift and rescale the boxes to fit in the output boundary.
The output boundary of the boxes can be smaller than the original image size
for accelerating the downstream calculations (dynamic mask resizing, mask IoU,
etc.).
For each image of the batch:
(1) find the upper boundary (min_ymin) and the left boundary (min_xmin) of all
the boxes.
(2) shift all the boxes up min_ymin pixels and left min_xmin pixels.
(3) find the new lower boundary (max_ymax) and the right boundary (max_xmax)
of all the boxes.
(4) if max_ymax > output_height or max_xmax > output_width (some boxes don't
fit in the output boundary), downsample all the boxes by ratio:
min(output_height / max_ymax, output_width / max_xmax). The aspect ratio
is not changed.
Args:
boxes: a tensor with a shape of [batch_size, N, 4]. The last dimension is
the pixel coordinates in [ymin, xmin, ymax, xmax] form.
output_boundary: two integers that represent the height and width of the
output.
Returns:
The tensor [batch_size, N, 4] of the output boxes.
"""
boxes = tf.cast(boxes, dtype=tf.float32)
# (batch_size, num_boxes, 1)
is_valid_box = tf.reduce_any(
(boxes[:, :, 2:4] - boxes[:, :, 0:2]) > 0, axis=-1, keepdims=True
)
# (batch_size, 2)
min_ymin_xmin = tf.reduce_min(
tf.where(is_valid_box, boxes, np.inf)[:, :, 0:2],
axis=1,
)
# (batch_size, num_boxes, 4)
boxes = tf.where(
is_valid_box,
boxes - tf.tile(min_ymin_xmin, [1, 2])[:, tf.newaxis, :],
0.0,
)
# (batch_size,)
max_ymax = tf.reduce_max(boxes[:, :, 2], axis=1)
max_xmax = tf.reduce_max(boxes[:, :, 3], axis=1)
# (batch_size,)
y_resize_ratio = output_boundary[0] / max_ymax
x_resize_ratio = output_boundary[1] / max_xmax
# (batch_size,)
downsampling_ratio = tf.math.minimum(
tf.math.minimum(y_resize_ratio, x_resize_ratio), 1.0
)
# (batch_size, num_boxes, 4)
return boxes * downsampling_ratio[:, tf.newaxis, tf.newaxis]
def _count_detection_type(
detection_type_mask: tf.Tensor,
detection_classes: tf.Tensor,
flattened_binned_confidence_one_hot: tf.Tensor,
num_classes: int,
) -> tf.Tensor:
"""Counts detection type grouped by IoU thresholds, classes and confidence bins.
Args:
detection_type_mask: a bool tensor in shape of (batch_size, num_detections,
num_iou_thresholds), which indicate a certain type of detections (e.g.
true postives).
detection_classes: a tensor in shape of (batch_size, num_detections) which
stores the classes of the detections.
flattened_binned_confidence_one_hot: a one-hot bool tensor in shape of
(batch_size * num_detections, num_confidence_bins + 1) which indicates the
binned confidence score of each detection.
num_classes: the number of classes.
Returns:
A tensor in shape of (num_iou_thresholds, num_classes,
num_confidence_bins + 1) which stores the count grouped by IoU thresholds,
classes and confidence bins.
"""
num_iou_thresholds = detection_type_mask.get_shape()[-1]
# (batch_size, num_detections, num_iou_thresholds)
masked_classes = tf.where(
detection_type_mask, detection_classes[..., tf.newaxis], -1
)
# (num_iou_thresholds, batch_size * num_detections)
flattened_masked_classes = tf.transpose(
tf.reshape(masked_classes, [-1, num_iou_thresholds])
)
# (num_iou_thresholds, num_classes, batch_size * num_detections)
flattened_masked_classes_one_hot = tf.one_hot(
flattened_masked_classes, depth=num_classes, axis=1
)
# (num_iou_thresholds * num_classes, batch_size * num_detections)
flattened_masked_classes_one_hot = tf.reshape(
flattened_masked_classes_one_hot,
[num_iou_thresholds * num_classes, -1],
)
# (num_iou_thresholds * num_classes, num_confidence_bins + 1)
count = tf.matmul(
flattened_masked_classes_one_hot,
tf.cast(flattened_binned_confidence_one_hot, tf.float32),
a_is_sparse=True,
b_is_sparse=True,
)
# (num_iou_thresholds, num_classes, num_confidence_bins + 1)
count = tf.reshape(count, [num_iou_thresholds, num_classes, -1])
# Clears the count of class 0 (background)
count *= 1.0 - tf.eye(num_classes, 1, dtype=count.dtype)
return count
class InstanceMetrics(tf_keras.metrics.Metric):
"""Reports the metrics of instance detection & segmentation."""
def __init__(
self,
num_classes: int,
use_masks: bool = False,
iou_thresholds: Tuple[float, ...] = (0.5,),
confidence_thresholds: Tuple[float, ...] = (),
num_confidence_bins: int = 1000,
mask_output_boundary: Tuple[int, int] = (640, 640),
matching_algorithm: Optional[MatchingAlgorithm] = None,
average_precision_algorithms: Optional[
Dict[str, AveragePrecision]
] = None,
name: Optional[str] = None,
dtype: Optional[Union[str, tf.dtypes.DType]] = tf.float32,
**kwargs
):
"""Initialization for AveragePrecision.
Args:
num_classes: the number of classes.
use_masks: if true, use the masks of the instances when calculating the
metrics, otherwise use the boxes.
iou_thresholds: a sequence of IoU thresholds over which to calculate the
instance metrics.
confidence_thresholds: a sequence of confidence thresholds. If set, also
report precision and recall at each confidence threshold, otherwise,
only report average precision.
num_confidence_bins: the number of confidence bins used for bin sort.
mask_output_boundary: two integers that represent the height and width of
the boundary where the resized instance masks are pasted. For each
example, if any of the detection or ground truth boxes is out of the
boundary, shift and resize all the detection and ground truth boxes of
the example to fit them into the boundary. The output boundary of the
pasted masks can be smaller than the real image size for accelerating
the calculation.
matching_algorithm: the algorithm which matches detections to ground
truths.
average_precision_algorithms: the algorithms which compute average
precision from P-R curve. The keys are used in the metrics results.
name: the name of the metric instance.
dtype: data type of the metric result.
**kwargs: Additional keywords arguments.
"""
super().__init__(name=name, dtype=dtype, **kwargs)
self._num_classes = num_classes
self._use_masks = use_masks
self._iou_thresholds = iou_thresholds
self._confidence_thresholds = confidence_thresholds
self._num_iou_thresholds = len(iou_thresholds)
self._num_confidence_bins = num_confidence_bins
self._mask_output_boundary = mask_output_boundary
if not matching_algorithm:
self._matching_algorithm = COCOMatchingAlgorithm(iou_thresholds)
else:
self._matching_algorithm = matching_algorithm
if not average_precision_algorithms:
self._average_precision_algorithms = {'ap': COCOAveragePrecision()}
else:
self._average_precision_algorithms = average_precision_algorithms
# Variables
self.tp_count = self.add_weight(
'tp_count',
shape=[
self._num_iou_thresholds,
self._num_classes,
self._num_confidence_bins + 1,
],
initializer='zeros',
dtype=tf.float32,
)
self.fp_count = self.add_weight(
'fp_count',
shape=[
self._num_iou_thresholds,
self._num_classes,
self._num_confidence_bins + 1,
],
initializer='zeros',
dtype=tf.float32,
)
self.gt_count = self.add_weight(
'gt_count',
shape=[self._num_classes],
initializer='zeros',
dtype=tf.float32,
)
def get_config(self) -> Dict[str, Any]:
"""Returns the serializable config of the metric."""
return {
'num_classes': self._num_classes,
'use_masks': self._use_masks,
'iou_thresholds': self._iou_thresholds,
'confidence_thresholds': self._confidence_thresholds,
'num_confidence_bins': self._num_confidence_bins,
'mask_output_boundary': self._mask_output_boundary,
'matching_algorithm': self._matching_algorithm,
'average_precision_algorithms': self._average_precision_algorithms,
'name': self.name,
'dtype': self.dtype,
}
def reset_state(self):
"""Resets all of the metric state variables."""
self.tp_count.assign(tf.zeros_like(self.tp_count))
self.fp_count.assign(tf.zeros_like(self.fp_count))
self.gt_count.assign(tf.zeros_like(self.gt_count))
def update_state(
self, y_true: Dict[str, tf.Tensor], y_pred: Dict[str, tf.Tensor]
):
# (batch_size, num_detections, 4) in absolute coordinates.
detection_boxes = tf.cast(y_pred['detection_boxes'], tf.float32)
# (batch_size, num_detections)
detection_classes = tf.cast(y_pred['detection_classes'], tf.int32)
# (batch_size, num_detections)
detection_scores = tf.cast(y_pred['detection_scores'], tf.float32)
# (batch_size, num_gts, 4) in absolute coordinates.
gt_boxes = tf.cast(y_true['boxes'], tf.float32)
# (batch_size, num_gts)
gt_classes = tf.cast(y_true['classes'], tf.int32)
# (batch_size, num_gts)
if 'is_crowds' in y_true:
gt_is_crowd = tf.cast(y_true['is_crowds'], tf.bool)
else:
gt_is_crowd = tf.zeros_like(gt_classes, dtype=tf.bool)
image_scale = tf.tile(y_true['image_info'][:, 2:3, :], multiples=[1, 1, 2])
detection_boxes = detection_boxes / tf.cast(
image_scale, dtype=detection_boxes.dtype
)
# Step 1: Computes IoUs between the detections and the non-crowd ground
# truths and IoAs between the detections and the crowd ground truths.
if not self._use_masks:
# (batch_size, num_detections, num_gts)
detection_to_gt_ious = box_ops.bbox_overlap(detection_boxes, gt_boxes)
detection_to_gt_ioas = box_ops.bbox_intersection_over_area(
detection_boxes, gt_boxes
)
else:
# Use outer boxes to generate the masks if available.
if 'detection_outer_boxes' in y_pred:
detection_boxes = tf.cast(y_pred['detection_outer_boxes'], tf.float32)
# (batch_size, num_detections, mask_height, mask_width)
detection_masks = tf.cast(y_pred['detection_masks'], tf.float32)
# (batch_size, num_gts, gt_mask_height, gt_mask_width)
gt_masks = tf.cast(y_true['masks'], tf.float32)
num_detections = detection_boxes.get_shape()[1]
# (batch_size, num_detections + num_gts, 4)
all_boxes = _shift_and_rescale_boxes(
tf.concat([detection_boxes, gt_boxes], axis=1),
self._mask_output_boundary,
)
detection_boxes = all_boxes[:, :num_detections, :]
gt_boxes = all_boxes[:, num_detections:, :]
# (batch_size, num_detections, num_gts)
detection_to_gt_ious, detection_to_gt_ioas = (
mask_ops.instance_masks_overlap(
detection_boxes,
detection_masks,
gt_boxes,
gt_masks,
output_size=self._mask_output_boundary,
)
)
# (batch_size, num_detections, num_gts)
detection_to_gt_ious = tf.where(
gt_is_crowd[:, tf.newaxis, :], 0.0, detection_to_gt_ious
)
detection_to_crowd_ioas = tf.where(
gt_is_crowd[:, tf.newaxis, :], detection_to_gt_ioas, 0.0
)
# Step 2: counts true positives grouped by IoU thresholds, classes and
# confidence bins.
# (batch_size, num_detections, num_iou_thresholds)
detection_is_tp, _ = self._matching_algorithm(
detection_to_gt_ious, detection_classes, detection_scores, gt_classes
)
# (batch_size * num_detections,)
flattened_binned_confidence = tf.reshape(
tf.cast(detection_scores * self._num_confidence_bins, tf.int32), [-1]
)
# (batch_size * num_detections, num_confidence_bins + 1)
flattened_binned_confidence_one_hot = tf.one_hot(
flattened_binned_confidence, self._num_confidence_bins + 1, axis=1
)
# (num_iou_thresholds, num_classes, num_confidence_bins + 1)
tp_count = _count_detection_type(
detection_is_tp,
detection_classes,
flattened_binned_confidence_one_hot,
self._num_classes,
)
# Step 3: Counts false positives grouped by IoU thresholds, classes and
# confidence bins.
# False positive: detection is not true positive (see above) and not part of
# the crowd ground truth with the same class.
# (batch_size, num_detections, num_gts, num_iou_thresholds)
detection_matches_crowd = (
(detection_to_crowd_ioas[..., tf.newaxis] > self._iou_thresholds)
& (
detection_classes[:, :, tf.newaxis, tf.newaxis]
== gt_classes[:, tf.newaxis, :, tf.newaxis]
)
& (detection_classes[:, :, tf.newaxis, tf.newaxis] > 0)
)
# (batch_size, num_detections, num_iou_thresholds)
detection_matches_any_crowd = tf.reduce_any(
detection_matches_crowd & ~detection_is_tp[:, :, tf.newaxis, :], axis=2
)
detection_is_fp = ~detection_is_tp & ~detection_matches_any_crowd
# (num_iou_thresholds, num_classes, num_confidence_bins + 1)
fp_count = _count_detection_type(
detection_is_fp,
detection_classes,
flattened_binned_confidence_one_hot,
self._num_classes,
)
# Step 4: Counts non-crowd groundtruths grouped by classes.
# (num_classes, )
gt_count = tf.reduce_sum(
tf.one_hot(
tf.where(gt_is_crowd, -1, gt_classes), self._num_classes, axis=-1
),
axis=[0, 1],
)
# Clears the count of class 0 (background).
gt_count *= 1.0 - tf.eye(1, self._num_classes, dtype=gt_count.dtype)[0]
# Accumulates the variables.
self.fp_count.assign_add(tf.cast(fp_count, self.fp_count.dtype))
self.tp_count.assign_add(tf.cast(tp_count, self.tp_count.dtype))
self.gt_count.assign_add(tf.cast(gt_count, self.gt_count.dtype))
def result(self) -> Dict[str, tf.Tensor]:
"""Returns the metrics values as a dict.
Returns:
A `dict` containing:
'ap': a float tensor in shape (num_iou_thresholds, num_classes) which
stores the average precision of each class at different IoU thresholds.
'precision': a float tensor in shape (num_confidence_thresholds,
num_iou_thresholds, num_classes) which stores the precision of each
class at different confidence thresholds & IoU thresholds.
'recall': a float tensor in shape (num_confidence_thresholds,
num_iou_thresholds, num_classes) which stores the recall of each
class at different confidence thresholds & IoU thresholds.
'valid_classes': a bool tensor in shape (num_classes,). If False, there
is no instance of the class in the ground truth.
"""
result = {
# (num_classes,)
'valid_classes': self.gt_count != 0,
}
# (num_iou_thresholds, num_classes, num_confidence_bins + 1)
tp_count_cum_by_confidence = tf.math.cumsum(
self.tp_count, axis=-1, reverse=True
)
# (num_iou_thresholds, num_classes, num_confidence_bins + 1)
fp_count_cum_by_confidence = tf.math.cumsum(
self.fp_count, axis=-1, reverse=True
)
# (num_iou_thresholds, num_classes, num_confidence_bins + 1)
precisions = tf.math.divide_no_nan(
tp_count_cum_by_confidence,
tp_count_cum_by_confidence + fp_count_cum_by_confidence,
)
# (num_iou_thresholds, num_classes, num_confidence_bins + 1)
recalls = tf.math.divide_no_nan(
tp_count_cum_by_confidence, self.gt_count[..., tf.newaxis]
)
if self._confidence_thresholds:
# If confidence_thresholds is set, reports precision and recall at each
# confidence threshold.
confidence_thresholds = tf.cast(
tf.constant(self._confidence_thresholds, dtype=tf.float32)
* self._num_confidence_bins,
dtype=tf.int32,
)
# (num_confidence_thresholds, num_iou_thresholds, num_classes)
result['precisions'] = tf.gather(
tf.transpose(precisions, [2, 0, 1]), confidence_thresholds
)
result['recalls'] = tf.gather(
tf.transpose(recalls, [2, 0, 1]), confidence_thresholds
)
precisions = tf.reverse(precisions, axis=[-1])
recalls = tf.reverse(recalls, axis=[-1])
result.update(
{
# (num_iou_thresholds, num_classes)
key: ap_algorithm(precisions, recalls)
for key, ap_algorithm in self._average_precision_algorithms.items()
}
)
return result
def get_average_precision_metrics_keys(self):
"""Gets the keys of the average precision metrics in the results."""
return self._average_precision_algorithms.keys()