# ----------------------------------------------------------------------------
# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
#
# 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.
# ----------------------------------------------------------------------------
import torch
[docs]
def batched_mask_to_box(masks: torch.Tensor) -> torch.Tensor:
"""Convert binary masks to bounding boxes.
Args:
masks (torch.Tensor): A tensor of shape (..., H, W) containing binary masks
where non-zero values indicate the presence of the object.
Returns:
torch.Tensor: A tensor of shape (..., 4) containing the bounding boxes
in XYXY format.
"""
# torch.max below raises an error on empty inputs, just skip in this case
if torch.numel(masks) == 0:
return torch.zeros(*masks.shape[:-2], 4, device=masks.device)
# Normalize shape to CxHxW
shape = masks.shape
h, w = shape[-2:]
if len(shape) > 2:
masks = masks.flatten(0, -3)
else:
masks = masks.unsqueeze(0)
# Get top and bottom edges
in_height, _ = torch.max(masks, dim=-1)
in_height_coords = in_height * torch.arange(h, device=in_height.device)[None, :]
bottom_edges, _ = torch.max(in_height_coords, dim=-1)
in_height_coords = in_height_coords + h * (~in_height)
top_edges, _ = torch.min(in_height_coords, dim=-1)
# Get left and right edges
in_width, _ = torch.max(masks, dim=-2)
in_width_coords = in_width * torch.arange(w, device=in_width.device)[None, :]
right_edges, _ = torch.max(in_width_coords, dim=-1)
in_width_coords = in_width_coords + w * (~in_width)
left_edges, _ = torch.min(in_width_coords, dim=-1)
# If the mask is empty the right edge will be to the left of the left edge.
# Replace these boxes with [0, 0, 0, 0]
empty_filter = (right_edges < left_edges) | (bottom_edges < top_edges)
out = torch.stack([left_edges, top_edges, right_edges, bottom_edges], dim=-1)
out = out * (~empty_filter).unsqueeze(-1)
# Return to original shape
if len(shape) > 2:
out = out.reshape(*shape[:-2], 4)
else:
out = out[0]
return out
[docs]
def gen_disp_colormap(inputs, normalize=True, torch_transpose=True):
"""Generate a color visualization with the ``plasma`` colormap.
Args:
inputs: NumPy array or tensor with shape ``(H, W)``, ``(N, H, W)``,
or ``(N, C, H, W)``. Four-dimensional inputs use the first channel.
normalize: Whether to scale the input linearly to ``[0, 1]`` before
applying the colormap. Defaults to True.
torch_transpose: Whether to transpose the generated array to
channel-first order. Defaults to True.
Returns:
A NumPy array containing the colormapped values. Its shape depends on
the input dimensionality and ``torch_transpose``.
Note:
Tensor inputs are detached and moved to CPU. The returned array is a
new value and does not modify the input.
"""
import matplotlib.pyplot as plt
import torch
_DEPTH_COLORMAP = plt.get_cmap("plasma", 256) # for plotting
if isinstance(inputs, torch.Tensor):
inputs = inputs.detach().cpu().numpy()
vis = inputs
if normalize:
ma = float(vis.max())
mi = float(vis.min())
d = ma - mi if ma != mi else 1e5
vis = (vis - mi) / d
if vis.ndim == 4:
vis = vis.transpose([0, 2, 3, 1])
vis = _DEPTH_COLORMAP(vis)
vis = vis[:, :, :, 0, :3]
if torch_transpose:
vis = vis.transpose(0, 3, 1, 2)
elif vis.ndim == 3:
vis = _DEPTH_COLORMAP(vis)
vis = vis[:, :, :3]
if torch_transpose:
vis = vis.transpose(0, 3, 1, 2)
elif vis.ndim == 2:
vis = _DEPTH_COLORMAP(vis)
vis = vis[..., :3]
if torch_transpose:
vis = vis.transpose(2, 0, 1)
return vis