official/vision/evaluation/segmentation_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 segmentation."""
from typing import Optional, Sequence, Tuple, Union
import tensorflow as tf, tf_keras
from official.vision.evaluation import iou
from official.vision.ops import box_ops
from official.vision.ops import spatial_transform_ops
class MeanIoU(tf_keras.metrics.MeanIoU):
"""Mean IoU metric for semantic segmentation.
This class utilizes tf_keras.metrics.MeanIoU to perform batched mean iou when
both input images and ground-truth masks are resized to the same size
(rescale_predictions=False). It also computes mean IoU on ground-truth
original sizes, in which case, each prediction is rescaled back to the
original image size.
"""
def __init__(self,
num_classes,
rescale_predictions=False,
name=None,
dtype=None):
"""Constructs Segmentation evaluator class.
Args:
num_classes: `int`, number of classes.
rescale_predictions: `bool`, whether to scale back prediction to original
image sizes. If True, y_true['image_info'] is used to rescale
predictions.
name: `str`, name of the metric instance..
dtype: data type of the metric result.
"""
self._rescale_predictions = rescale_predictions
super().__init__(num_classes=num_classes, name=name, dtype=dtype)
def update_state(self, y_true, y_pred):
"""Updates metric state.
Args:
y_true: `dict`, dictionary with the following name, and key values.
- masks: [batch, height, width, 1], ground-truth masks.
- valid_masks: [batch, height, width, 1], valid elements in the mask.
- image_info: [batch, 4, 2], a tensor that holds information about
original and preprocessed images. Each entry is in the format of
[[original_height, original_width], [input_height, input_width],
[y_scale, x_scale], [y_offset, x_offset]], where [desired_height,
desired_width] is the actual scaled image size, and [y_scale, x_scale]
is the scaling factor, which is the ratio of scaled dimension /
original dimension.
y_pred: Tensor [batch, height_p, width_p, num_classes], predicated masks.
"""
predictions, masks, valid_masks = preprocess_inputs(
y_true, y_pred, self._rescale_predictions)
# Ignored mask elements are set to zero for fitting the confusion matrix.
masks = tf.where(valid_masks, masks, tf.zeros_like(masks))
predictions = tf.argmax(predictions, axis=3)
flatten_predictions = tf.reshape(predictions, shape=[-1])
flatten_masks = tf.reshape(masks, shape=[-1])
flatten_valid_masks = tf.reshape(valid_masks, shape=[-1])
super().update_state(
y_true=flatten_masks,
y_pred=flatten_predictions,
sample_weight=tf.cast(flatten_valid_masks, tf.float32))
class PerClassIoU(MeanIoU):
"""Per class IoU metric for semantic segmentation."""
def result(self):
"""Compute IoU for each class via the confusion matrix."""
sum_over_row = tf.cast(
tf.reduce_sum(self.total_cm, axis=0), dtype=self._dtype)
sum_over_col = tf.cast(
tf.reduce_sum(self.total_cm, axis=1), dtype=self._dtype)
true_positives = tf.cast(
tf.linalg.tensor_diag_part(self.total_cm), dtype=self._dtype)
# sum_over_row + sum_over_col =
# 2 * true_positives + false_positives + false_negatives.
denominator = sum_over_row + sum_over_col - true_positives
return tf.math.divide_no_nan(true_positives, denominator)
class PerClassIoUV2(iou.PerClassIoUV2):
"""Computes the per-class IoU metric for semantic segmentation.
This implementation converts predictions and ground truth to binary masks,
and uses logical AND and OR to compute intersection and union, which is much
faster than the MeanIoU and PerClassIoU (using confusion matrix) above on TPU,
but slower on CPU and GPU.
"""
def __init__(self,
num_classes: int,
rescale_predictions: bool = False,
name: Optional[str] = None,
dtype: Optional[Union[str, tf.dtypes.DType]] = tf.float32,
shape: Optional[Sequence[int]] = None,
axis: int = -1):
"""Constructs Segmentation evaluator class.
Args:
num_classes: `int`, number of classes.
rescale_predictions: `bool`, whether to scale back prediction to original
image sizes. If True, y_true['image_info'] is used to rescale
predictions.
name: `str`, name of the metric instance.
dtype: data type of the metric result.
shape: shape of the metrics result.
axis: (Optional) Defaults to -1. The dimension containing the one-hot
values.
"""
super().__init__(
num_classes=num_classes, name=name, dtype=dtype, shape=shape, axis=axis)
self._rescale_predictions = rescale_predictions
def update_state(self, y_true: tf.Tensor, y_pred: tf.Tensor):
"""Updates metric state.
Args:
y_true: `dict`, dictionary with the following name, and key values.
- masks: [batch, height, width, num_layers], ground-truth masks. The
num_layers is 1 by default, while all the operations in this function
support num_layers > 1.
- valid_masks: [batch, height, width, num_layers], valid elements in the
mask.
- image_info: [batch, 4, 2], a tensor that holds information about
original and preprocessed images. Each entry is in the format of
[[original_height, original_width], [input_height, input_width],
[y_scale, x_scale], [y_offset, x_offset]], where [desired_height,
desired_width] is the actual scaled image size, and [y_scale, x_scale]
is the scaling factor, which is the ratio of scaled dimension /
original dimension.
y_pred: Tensor [batch, height_p, width_p, num_classes], predicated masks.
"""
logits, gt_masks, valid_masks = preprocess_inputs(y_true, y_pred,
self._rescale_predictions)
valid_masks = tf.cast(valid_masks, tf.bool)
gt_binary_masks = tf.one_hot(
tf.cast(gt_masks[..., 0], dtype=tf.int32),
depth=self.num_classes,
on_value=True,
off_value=False,
)
gt_binary_masks &= valid_masks
predictions_binary_masks = tf.one_hot(
tf.argmax(logits, axis=-1, output_type=tf.int32),
depth=self.num_classes,
on_value=True,
off_value=False,
)
predictions_binary_masks &= valid_masks
super().update_state(
y_true=gt_binary_masks, y_pred=predictions_binary_masks
)
class MeanIoUV2(PerClassIoUV2):
"""Computes the mean IoU metric for semantic segmentation."""
def __init__(self,
target_class_ids: Optional[Tuple[int, ...]] = None,
**kwargs):
"""Initializes the class.
Args:
target_class_ids: computes mean IoU for the target classes. Selects all
the if empty.
**kwargs: the other arguments for initializing the base class.
"""
super().__init__(**kwargs)
self._target_class_ids = target_class_ids
def result(self) -> tf.Tensor:
"""Average the IoUs of all the classes."""
# (num_classes, )
per_class_ious = super().result()
if self._target_class_ids:
# (num_classes, )
target_class_indicators = tf.reduce_max(
tf.one_hot(
self._target_class_ids,
depth=self.num_classes,
dtype=per_class_ious.dtype),
axis=0)
return tf.math.divide_no_nan(
tf.reduce_sum(per_class_ious * target_class_indicators),
tf.reduce_sum(target_class_indicators))
else:
return tf.reduce_mean(per_class_ious)
def preprocess_inputs(
y_true: tf.Tensor, y_pred: tf.Tensor,
rescale_predictions: bool) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
"""Pre-processes the inputs (predictions and ground-truth) of the metrics.
Args:
y_true: `dict`, dictionary with the following name, and key values.
- masks: [batch, height, width, num_layers], ground-truth masks. The
num_layers is 1 by default, while all the operations in this function
support num_layers > 1.
- valid_masks: [batch, height, width, num_layers], valid elements in the
mask.
- image_info: [batch, 4, 2], a tensor that holds information about
original and preprocessed images. Each entry is in the format of
[[original_height, original_width], [input_height, input_width],
[y_scale, x_scale], [y_offset, x_offset]], where [desired_height,
desired_width] is the actual scaled image size, and [y_scale, x_scale]
is the scaling factor, which is the ratio of scaled dimension /
original dimension.
y_pred: tensor [batch, height_p, width_p, num_classes], predicated masks.
rescale_predictions: `bool`, whether to scale back prediction to original
image sizes. If True, y_true['image_info'] is used to rescale predictions.
Returns:
logits: a float tensor in shape [batch, height, width, num_classes], which
stores the raw output of the model.
gt_masks: an int tensor in shape [batch, height, width, 1], which stores the
ground-truth masks.
valid_masks: a bool tensor in shape [batch, height, width, 1], which
indicates the valid elements of the masks.
"""
logits = y_pred
gt_masks = y_true['masks']
valid_masks = y_true['valid_masks']
images_info = y_true['image_info']
if isinstance(logits, tuple) or isinstance(logits, list):
logits = tf.concat(logits, axis=0)
gt_masks = tf.concat(gt_masks, axis=0)
valid_masks = tf.concat(valid_masks, axis=0)
images_info = tf.concat(images_info, axis=0)
# The pixel is valid if any layer of the masks is valid at that pixel.
# (batch_size, height, width)
valid_masks = tf.reduce_any(tf.cast(valid_masks, tf.bool), axis=-1)
gt_masks_size = tf.shape(gt_masks)[1:3]
if rescale_predictions:
# Scale back predictions to original image shapes and pad to mask size.
# Note: instead of cropping the masks to image shape (dynamic), here we
# pad the rescaled predictions to mask size (fixed). And update the
# valid_masks to mask out the pixels outside the original image shape.
logits, image_shape_masks = (
_rescale_and_pad_predictions(
logits, images_info, output_size=gt_masks_size))
# Only the area within the original image shape is valid.
# (batch_size, height, width)
valid_masks &= image_shape_masks
else:
logits = tf.image.resize(
logits, gt_masks_size, method=tf.image.ResizeMethod.BILINEAR)
# (batch_size, height, width, 1)
valid_masks = valid_masks[..., tf.newaxis]
return logits, gt_masks, valid_masks
def _rescale_and_pad_predictions(
predictions: tf.Tensor, images_info: tf.Tensor,
output_size: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
"""Scales back predictions to original image shapes and pads to output size.
Args:
predictions: A tensor in shape [batch, height, width, num_classes] which
stores the model predictions.
images_info: A tensor in shape [batch, 4, 2] that holds information about
original and preprocessed images. Each entry is in the format of
[[original_height, original_width], [input_height, input_width], [y_scale,
x_scale], [y_offset, x_offset]], where [desired_height, desired_width] is
the actual scaled image size, and [y_scale, x_scale] is the scaling
factor, which is the ratio of scaled dimension / original dimension.
output_size: A list/tuple/tensor stores the size of the padded output in
[output_height, output_width].
Returns:
predictions: A tensor in shape [batch, output_height, output_width,
num_classes] which stores the rescaled and padded predictions.
image_shape_masks: A bool tensor in shape [batch, output_height,
output_width] where the pixels inside the original image shape are true,
otherwise false.
"""
# (batch_size, 2)
image_shape = tf.cast(images_info[:, 0, :], tf.int32)
desired_size = tf.cast(images_info[:, 1, :], tf.float32)
image_scale = tf.cast(images_info[:, 2, :], tf.float32)
offset = tf.cast(images_info[:, 3, :], tf.int32)
rescale_size = tf.cast(tf.math.ceil(desired_size / image_scale), tf.int32)
# Rescale the predictions, then crop to the original image shape and
# finally pad zeros to match the mask size.
predictions = (
spatial_transform_ops.bilinear_resize_with_crop_and_pad(
predictions,
rescale_size,
crop_offset=offset,
crop_size=image_shape,
output_size=output_size))
# (batch_size, 2)
y0_x0 = tf.broadcast_to(
tf.constant([[0, 0]], dtype=image_shape.dtype), tf.shape(image_shape))
# (batch_size, 4)
image_shape_bbox = tf.concat([y0_x0, image_shape], axis=1)
# (batch_size, height, width)
image_shape_masks = box_ops.bbox2mask(
bbox=image_shape_bbox,
image_height=output_size[0],
image_width=output_size[1],
dtype=tf.bool)
return predictions, image_shape_masks