official/vision/ops/mask_ops.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.
"""Utility functions for segmentations."""
import math
from typing import List, Tuple
# Import libraries
import cv2
import numpy as np
import tensorflow as tf, tf_keras
from official.vision.ops import spatial_transform_ops
def paste_instance_masks(masks: np.ndarray, detected_boxes: np.ndarray,
image_height: int, image_width: int) -> np.ndarray:
"""Paste instance masks to generate the image segmentation results.
Args:
masks: a numpy array of shape [N, mask_height, mask_width] representing the
instance masks w.r.t. the `detected_boxes`.
detected_boxes: a numpy array of shape [N, 4] representing the reference
bounding boxes.
image_height: an integer representing the height of the image.
image_width: an integer representing the width of the image.
Returns:
segms: a numpy array of shape [N, image_height, image_width] representing
the instance masks *pasted* on the image canvas.
"""
def expand_boxes(boxes: np.ndarray, scale: float) -> np.ndarray:
"""Expands an array of boxes by a given scale."""
# Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/utils/boxes.py#L227 # pylint: disable=line-too-long
# The `boxes` in the reference implementation is in [x1, y1, x2, y2] form,
# whereas `boxes` here is in [x1, y1, w, h] form
w_half = boxes[:, 2] * 0.5
h_half = boxes[:, 3] * 0.5
x_c = boxes[:, 0] + w_half
y_c = boxes[:, 1] + h_half
w_half *= scale
h_half *= scale
boxes_exp = np.zeros(boxes.shape)
boxes_exp[:, 0] = x_c - w_half
boxes_exp[:, 2] = x_c + w_half
boxes_exp[:, 1] = y_c - h_half
boxes_exp[:, 3] = y_c + h_half
return boxes_exp
# Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/core/test.py#L812 # pylint: disable=line-too-long
# To work around an issue with cv2.resize (it seems to automatically pad
# with repeated border values), we manually zero-pad the masks by 1 pixel
# prior to resizing back to the original image resolution. This prevents
# "top hat" artifacts. We therefore need to expand the reference boxes by an
# appropriate factor.
_, mask_height, mask_width = masks.shape
scale = max((mask_width + 2.0) / mask_width,
(mask_height + 2.0) / mask_height)
ref_boxes = expand_boxes(detected_boxes, scale)
ref_boxes = ref_boxes.astype(np.int32)
padded_mask = np.zeros((mask_height + 2, mask_width + 2), dtype=np.float32)
segms = []
for mask_ind, mask in enumerate(masks):
im_mask = np.zeros((image_height, image_width), dtype=np.uint8)
# Process mask inside bounding boxes.
padded_mask[1:-1, 1:-1] = mask[:, :]
ref_box = ref_boxes[mask_ind, :]
w = ref_box[2] - ref_box[0] + 1
h = ref_box[3] - ref_box[1] + 1
w = np.maximum(w, 1)
h = np.maximum(h, 1)
mask = cv2.resize(padded_mask, (w, h))
mask = np.array(mask > 0.5, dtype=np.uint8)
x_0 = min(max(ref_box[0], 0), image_width)
x_1 = min(max(ref_box[2] + 1, 0), image_width)
y_0 = min(max(ref_box[1], 0), image_height)
y_1 = min(max(ref_box[3] + 1, 0), image_height)
im_mask[y_0:y_1, x_0:x_1] = mask[
(y_0 - ref_box[1]):(y_1 - ref_box[1]),
(x_0 - ref_box[0]):(x_1 - ref_box[0])
]
segms.append(im_mask)
segms = np.array(segms)
assert masks.shape[0] == segms.shape[0]
return segms
def paste_instance_masks_v2(masks: np.ndarray, detected_boxes: np.ndarray,
image_height: int, image_width: int) -> np.ndarray:
"""Paste instance masks to generate the image segmentation (v2).
Args:
masks: a numpy array of shape [N, mask_height, mask_width] representing the
instance masks w.r.t. the `detected_boxes`.
detected_boxes: a numpy array of shape [N, 4] representing the reference
bounding boxes.
image_height: an integer representing the height of the image.
image_width: an integer representing the width of the image.
Returns:
segms: a numpy array of shape [N, image_height, image_width] representing
the instance masks *pasted* on the image canvas.
"""
_, mask_height, mask_width = masks.shape
segms = []
for i, mask in enumerate(masks):
box = detected_boxes[i, :]
xmin = box[0]
ymin = box[1]
xmax = xmin + box[2]
ymax = ymin + box[3]
# Sample points of the cropped mask w.r.t. the image grid.
# Note that these coordinates may fall beyond the image.
# Pixel clipping will happen after warping.
xmin_int = int(math.floor(xmin))
xmax_int = int(math.ceil(xmax))
ymin_int = int(math.floor(ymin))
ymax_int = int(math.ceil(ymax))
alpha = box[2] / (1.0 * mask_width)
beta = box[3] / (1.0 * mask_height)
# pylint: disable=invalid-name
# Transformation from mask pixel indices to image coordinate.
M_mask_to_image = np.array(
[[alpha, 0, xmin],
[0, beta, ymin],
[0, 0, 1]],
dtype=np.float32)
# Transformation from image to cropped mask coordinate.
M_image_to_crop = np.array(
[[1, 0, -xmin_int],
[0, 1, -ymin_int],
[0, 0, 1]],
dtype=np.float32)
M = np.dot(M_image_to_crop, M_mask_to_image)
# Compensate the half pixel offset that OpenCV has in the
# warpPerspective implementation: the top-left pixel is sampled
# at (0,0), but we want it to be at (0.5, 0.5).
M = np.dot(
np.dot(
np.array([[1, 0, -0.5],
[0, 1, -0.5],
[0, 0, 1]], np.float32),
M),
np.array([[1, 0, 0.5],
[0, 1, 0.5],
[0, 0, 1]], np.float32))
# pylint: enable=invalid-name
cropped_mask = cv2.warpPerspective(
mask.astype(np.float32), M,
(xmax_int - xmin_int, ymax_int - ymin_int))
cropped_mask = np.array(cropped_mask > 0.5, dtype=np.uint8)
img_mask = np.zeros((image_height, image_width))
x0 = max(min(xmin_int, image_width), 0)
x1 = max(min(xmax_int, image_width), 0)
y0 = max(min(ymin_int, image_height), 0)
y1 = max(min(ymax_int, image_height), 0)
img_mask[y0:y1, x0:x1] = cropped_mask[
(y0 - ymin_int):(y1 - ymin_int),
(x0 - xmin_int):(x1 - xmin_int)]
segms.append(img_mask)
segms = np.array(segms)
return segms
def instance_masks_overlap(
boxes: tf.Tensor,
masks: tf.Tensor,
gt_boxes: tf.Tensor,
gt_masks: tf.Tensor,
output_size: List[int],
mask_binarize_threshold: float = 0.5,
) -> Tuple[tf.Tensor, tf.Tensor]:
"""Calculates the IoUs and IoAs between the detection masks and the ground truth masks.
IoU: intersection over union.
IoA: intersection over the area of the detection masks.
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.
masks: a float tensor with a shape of [batch_size, N, mask_height,
mask_width] representing the instance masks w.r.t. the `boxes`.
gt_boxes: a tensor with a shape of [batch_size, M, 4]. The last dimension is
the pixel coordinates in [ymin, xmin, ymax, xmax] form.
gt_masks: a float tensor with a shape of [batch_size, M, gt_mask_height,
gt_mask_width] representing the instance masks w.r.t. the `gt_boxes`.
output_size: two integers that represent the height and width of the output
masks.
mask_binarize_threshold: a float representing the threshold for binarizing
mask values. Default value is 0.5.
Returns:
iou: a tensor with as a shape of [batch_size, N, M].
"""
_, num_detections, mask_height, mask_width = masks.get_shape().as_list()
_, num_gts, gt_mask_height, gt_mask_width = gt_masks.get_shape().as_list()
output_height, output_width = output_size
masks = tf.where(masks < 0, tf.zeros_like(masks), masks)
gt_masks = tf.where(gt_masks < 0, tf.zeros_like(gt_masks), gt_masks)
pasted_masks = tf.reshape(
spatial_transform_ops.bilinear_resize_to_bbox(
tf.reshape(masks, [-1, mask_height, mask_width]),
tf.reshape(boxes, [-1, 4]),
output_size,
),
shape=[-1, num_detections, output_height, output_width],
)
pasted_gt_masks = tf.reshape(
spatial_transform_ops.bilinear_resize_to_bbox(
tf.reshape(gt_masks, [-1, gt_mask_height, gt_mask_width]),
tf.reshape(gt_boxes, [-1, 4]),
output_size,
),
shape=[-1, num_gts, output_height, output_width],
)
# (batch_size, num_detections, output_height * output_width)
flattened_binary_masks = tf.reshape(
pasted_masks > mask_binarize_threshold,
[-1, num_detections, output_height * output_width],
)
# (batch_size, num_gts, output_height * output_width)
flattened_gt_binary_masks = tf.reshape(
pasted_gt_masks > mask_binarize_threshold,
[-1, num_gts, output_height * output_width],
)
# (batch_size, output_height * output_width, num_gts)
flattened_gt_binary_masks = tf.transpose(flattened_gt_binary_masks, [0, 2, 1])
flattened_binary_masks = tf.cast(flattened_binary_masks, tf.float32)
flattened_gt_binary_masks = tf.cast(flattened_gt_binary_masks, tf.float32)
# (batch_size, num_detections, num_gts)
intersection = tf.matmul(flattened_binary_masks, flattened_gt_binary_masks)
detection_area = tf.reduce_sum(flattened_binary_masks, axis=-1, keepdims=True)
gt_area = tf.reduce_sum(flattened_gt_binary_masks, axis=-2, keepdims=True)
union = detection_area + gt_area - intersection
return tf.math.divide_no_nan(intersection, union), tf.math.divide_no_nan(
intersection, detection_area
)