feat: support for PyTorch tensors in cc3d.connected_components

automated_test.py

@@ -1421,6 +1421,25 @@ def test_color_connectivity_graph_26():
   cc_labels = cc3d.color_connectivity_graph(vcg, connectivity=26)
   assert np.all(cc_labels == 1)
 
+def test_pytorch_integration_ccl_doesnt_crash():
+    torch = pytest.importorskip("torch")
+
+    labels = torch.from_numpy(np.zeros([100,100,100], dtype=np.uint16))
+
+    out = cc3d.connected_components(labels)
+
+    assert isinstance(out, torch.Tensor)
+    assert torch.all(out == labels)
+
+
+
+
+
+
+
+
+
+
 
 
 

cc3d.pyx

@@ -293,6 +293,11 @@ def connected_components(
     if return_N: (OUT, N)
     else: OUT
   """
+  is_torch = hasattr(data, "cpu")
+  if is_torch:
+    # don't need to call .detach() b/c its read-only
+    data = data.cpu().numpy() 
+
   cdef int dims = len(data.shape)
   if dims not in (1,2,3):
     raise DimensionError("Only 1D, 2D, and 3D arrays supported. Got: " + str(dims))
@@ -569,6 +574,10 @@ def connected_components(
 
   out_labels = _final_reshape(out_labels, sx, sy, sz, dims, order)
 
+  if is_torch:
+    import torch
+    out_labels = torch.from_numpy(out_labels)
+
   if return_N:
     return (out_labels, N)
   return out_labels