File size: 2,870 Bytes
c6535db | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | # Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
import logging
import torch
try:
from cc_torch import get_connected_components
HAS_CC_TORCH = True
except ImportError:
logging.debug(
"cc_torch not found. Consider installing for better performance. Command line:"
" pip install git+https://github.com/ronghanghu/cc_torch.git"
)
HAS_CC_TORCH = False
def connected_components_cpu_single(values: torch.Tensor):
assert values.dim() == 2
from skimage.measure import label
labels, num = label(values.cpu().numpy(), return_num=True)
labels = torch.from_numpy(labels)
counts = torch.zeros_like(labels)
for i in range(1, num + 1):
cur_mask = labels == i
cur_count = cur_mask.sum()
counts[cur_mask] = cur_count
return labels, counts
def connected_components_cpu(input_tensor: torch.Tensor):
out_shape = input_tensor.shape
if input_tensor.dim() == 4 and input_tensor.shape[1] == 1:
input_tensor = input_tensor.squeeze(1)
else:
assert (
input_tensor.dim() == 3
), "Input tensor must be (B, H, W) or (B, 1, H, W)."
batch_size = input_tensor.shape[0]
labels_list = []
counts_list = []
for b in range(batch_size):
labels, counts = connected_components_cpu_single(input_tensor[b])
labels_list.append(labels)
counts_list.append(counts)
labels_tensor = torch.stack(labels_list, dim=0).to(input_tensor.device)
counts_tensor = torch.stack(counts_list, dim=0).to(input_tensor.device)
return labels_tensor.view(out_shape), counts_tensor.view(out_shape)
def connected_components(input_tensor: torch.Tensor):
"""
Computes connected components labeling on a batch of 2D tensors, using the best available backend.
Args:
input_tensor (torch.Tensor): A BxHxW integer tensor or Bx1xHxW. Non-zero values are considered foreground. Bool tensor also accepted
Returns:
Tuple[torch.Tensor, torch.Tensor]: Both tensors have the same shape as input_tensor.
- A tensor with dense labels. Background is 0.
- A tensor with the size of the connected component for each pixel.
"""
if input_tensor.dim() == 3:
input_tensor = input_tensor.unsqueeze(1)
assert (
input_tensor.dim() == 4 and input_tensor.shape[1] == 1
), "Input tensor must be (B, H, W) or (B, 1, H, W)."
if input_tensor.is_cuda:
if HAS_CC_TORCH:
return get_connected_components(input_tensor.to(torch.uint8))
else:
# triton fallback
from sam3.perflib.triton.connected_components import (
connected_components_triton,
)
return connected_components_triton(input_tensor)
# CPU fallback
return connected_components_cpu(input_tensor)
|