search

prittt / YACCLAB

YACCLAB: Yet Another Connected Components Labeling Benchmark
BSD 3-Clause "New" or "Revised" License
203 stars 37 forks source link

Connected Components with PyTorch #28

Closed helia95 closed 2 years ago

helia95 commented 3 years ago

Hi, thanks for this great resource. I'm wondering if, in the spirit of making this even more accessible, you plan to add the possibility to create torch CUDA extensions (https://pytorch.org/tutorials/advanced/cpp_extension.html).

Many thanks

CostantinoGrana commented 3 years ago

This could be done, and I don't see any real problem with it, but I believe that only one or two algorithms should be ported (the best performing ones). What would the use of CCL in PyTorch be? How would this be called/used? What output would be useful?

MicheleCancilla commented 3 years ago

The Kornia library includes an iterative CCL implementation (PR). Additionally, @zsef123 already wrapped one of the GPU YACCLAB algorithm to be used together with PyTorch in zsef123/Connected_components_PyTorch.

helia95 commented 2 years ago

Hi, sorry for the late reply.

@CostantinoGrana I think the use case is something like (a more complex) pooling layer (e.g. average pooling link). The CCL would be used as a nn.Module, that can be created using torch extension (link in the first comment). Given a batch of examples as input of shape [n, 1, H, W], I think would be useful to provide the output in two ways (rough ideas):

  1. a tensor of shape [n, 1, H, W] where locations in (H, W) that belong to the same connected components are assigned the same label (similar to OpenCV).
  2. for each example in n, return the list of indices in (H, W) that belong to the same connected component (in this way it would be easier to pool the elements from the input tensor).

@MicheleCancilla Thanks for pointing it out. I think the problems with these solutions are

  1. the algorithm implemented in Korina library doesn't seem as reliable/efficient as the ones implemented here (but correct me if I'm wrong).
  2. I tried the other solution but it is broken, or at least I couldn't make it work. In general, there is no support for it.

Thank you!

MarcoForte commented 2 years ago

@helia95 I managed to make it work. Using PyTorch 1.10 and moving the notebook to a different folder.

It is much faster than the connected components in Kornia right now. For the example image they use at 512x512. Their CC is 40.1 µs, kornia is 9.13 ms

I think it would be great to port to Kornia