Python:如何从3D numpy / torch数组中提取连接的组件(边界框)?

Question

我在NumPy/Torch中为3D数组设置了二进制分割掩码。我想将这些转换为边界框(又名连接组件)。作为免责声明，每个数组可以包含多个连接的组件/边界框，这意味着我不能只取最小和最大非零索引值。

具体来说，假设我有一个二进制值的3D数组(我将使用2D，因为2D更容易可视化)。我想知道连接的分量是什么。例如，我想采用这个分割掩码:

>>> segmentation_mask
array([[1, 0, 0, 0, 0],
[0, 1, 0, 0, 0],
[1, 1, 1, 0, 0],
[1, 1, 0, 1, 0],
[1, 1, 0, 0, 1]], dtype=int32)

并将其转换为连接的组件，其中连接的组件具有任意标签，即

>>> connected_components
array([[1, 0, 0, 0, 0],
[0, 2, 0, 0, 0],
[2, 2, 2, 0, 0],
[2, 2, 0, 3, 0],
[2, 2, 0, 0, 4]], dtype=int32)

我如何用3D数组做到这一点?我愿意使用Numpy, Scipy, Torchvision, opencv，任何库。

Answer 1

这应该适用于任何数量的维度:

import numpy as np                                                                
            
from scipy.sparse import csr_matrix                                               
from scipy.sparse.csgraph import connected_components                             
            
segmentation_mask = np.array([[1, 0, 0, 0, 0],                                    
[0, 1, 0, 0, 0],                                    
[1, 1, 1, 0, 0],                                    
[1, 1, 0, 1, 0],                                    
[1, 1, 0, 0, 1]], dtype=np.int32)                   
            
row = []                                                                          
col = []                                                                          
segmentation_mask_reader = segmentation_mask.reshape(-1)                          
n_nodes = len(segmentation_mask_reader)                                           
for node in range(n_nodes):                                                       
idxs = np.unravel_index(node, segmentation_mask.shape)                        
if segmentation_mask[idxs] == 0:                                              
col.append(n_nodes)                                                       
else:                                                                         
for i in range(len(idxs)):                                                
if idxs[i] > 0:                                                       
new_idxs = list(idxs)                                             
new_idxs[i] -= 1                                                  
new_node = np.ravel_multi_index(new_idxs, segmentation_mask.shape)
if segmentation_mask_reader[new_node] != 0:                       
col.append(new_node)                                          
while len(col) > len(row):                                                    
row.append(node)                                                          
            
row = np.array(row, dtype=np.int32)                                               
col = np.array(col, dtype=np.int32)                                               
data = np.ones(len(row), dtype=np.int32)                                          
            
graph = csr_matrix((np.array(data), (np.array(row), np.array(col))),              
shape=(n_nodes+1, n_nodes+1))                                  
n_components, labels = connected_components(csgraph=graph)                        
            
background_label = labels[-1]                                                     
solution = np.zeros(segmentation_mask.shape, dtype=segmentation_mask.dtype)       
solution_writer = solution.reshape(-1)                                            
for node in range(n_nodes):                                                       
label = labels[node]                                                          
if label < background_label:                                                  
solution_writer[node] = label+1                                           
elif label > background_label:                                                
solution_writer[node] = label                                             
            
print(solution)