Connected Components Labeling for PyTorch

Update notes 2024.12

• Fix bugs and test in PyTorch version 1.13.1+cu117 and 2.1.2+cu118
• Installation: python setup.py install

Update notes 2026.01

• Simplify code structure for c++ compilation

CMakeLists.txt examples

set(CC_TORCH_SRC_DIR ${PROJECT_SOURCE_DIR}/third_party/Connected_components_PyTorch/cpp)

add_library(cc_torch SHARED
    ${CC_TORCH_SRC_DIR}/buf.h
    ${CC_TORCH_SRC_DIR}/buf.cu)

set_target_properties(cc_torch PROPERTIES CUDA_ARCHITECTURES "75;86")

target_include_directories(cc_torch PUBLIC
    ${CC_TORCH_SRC_DIR}
    ${TORCH_INCLUDE_DIRS}
)

target_link_libraries(cc_torch PUBLIC
    ${TORCH_LIBRARIES}
)

Usage in c++

#include "third_party/Connected_components_PyTorch/cpp/buf.h"

auto labeled_mask = cc_torch::connected_components_labeling_2d(mask.to(torch::kUInt8));

References

• YACCLAB : Yet Another Connected Components Labeling Benchmark

• Allegretti, Stefano; Bolelli, Federico; Grana, Costantino "Optimized Block-Based Algorithms to Label Connected Components on GPUs." IEEE Transactions on Parallel and Distributed Systems, 2019. BibTex.

---

Follwing Block-based Union Find Algorithm from YACCLAB.

• Running on GPU.
• PyTorch Interface
• Fix some bit alignment problem
• little refactoring

Example

Tested on scikit-image example ,follows example.ipynb

Install

> python3 setup.py install

> python3 setup.py test