Merge pull request #13 from masadcv/cuda-multichannel-support

Add multi-channel support for CUDA distance transforms
masadcv · Jul 16, 2022 · ab7a599 · ab7a599
2 parents ac36e64 + b6f3b3d
commit ab7a599
Show file tree

Hide file tree

Showing 7 changed files with 429 additions and 355 deletions.
diff --git a/FastGeodis/common.h b/FastGeodis/common.h
@@ -54,12 +54,12 @@ void check_spatial_shape_match(const torch::Tensor &in1, const torch::Tensor &in
 {
     if (in1.dim() != in2.dim())
     {
-        throw std::invalid_argument("dimensions of input tensors do not match " 
-            + std::to_string(in1.dim() - 2) + " vs " + std::to_string(in2.dim() - 2));
+        throw std::invalid_argument("dimensions of input tensors do not match " + \
+            std::to_string(in1.dim() - 2) + " vs " + std::to_string(in2.dim() - 2));
     }
-    for(int i=0; i < dims; i++)
+    for (int i = 0; i < dims; i++)
     {
-        if(in1.size(2+i) != in2.size(2+i))
+        if (in1.size(2 + i) != in2.size(2 + i))
         {
             std::cout << "Tensor1 ";
             print_shape(in1);
@@ -116,5 +116,5 @@ void check_input_dimensions(const torch::Tensor &image, const torch::Tensor &mas
     check_single_batch(mask);
 
     // check spatial shapes match
-    check_spatial_shape_match(image, mask, num_dims-2);    
+    check_spatial_shape_match(image, mask, num_dims - 2);
 }
diff --git a/FastGeodis/fastgeodis.h b/FastGeodis/fastgeodis.h
@@ -36,57 +36,53 @@
 
 #ifdef WITH_CUDA
 torch::Tensor generalised_geodesic2d_cuda(
-    torch::Tensor &image, 
-    const torch::Tensor &mask, 
-    const float &v, 
-    const float &l_grad, 
-    const float &l_eucl, 
+    torch::Tensor &image,
+    const torch::Tensor &mask,
+    const float &v,
+    const float &l_grad,
+    const float &l_eucl,
     const int &iterations);
 
 torch::Tensor generalised_geodesic3d_cuda(
-    torch::Tensor &image, 
-    const torch::Tensor &mask, 
-    const std::vector<float> &spacing, 
-    const float &v, 
-    const float &l_grad, 
-    const float &l_eucl, 
+    torch::Tensor &image,
+    const torch::Tensor &mask,
+    const std::vector<float> &spacing,
+    const float &v,
+    const float &l_grad,
+    const float &l_eucl,
     const int &iterations);
 #endif
 
 torch::Tensor generalised_geodesic2d_cpu(
-    torch::Tensor &image, 
-    const torch::Tensor &mask, 
-    const float &v, 
-    const float &l_grad, 
-    const float &l_eucl, 
+    torch::Tensor &image,
+    const torch::Tensor &mask,
+    const float &v,
+    const float &l_grad,
+    const float &l_eucl,
     const int &iterations);
 
 torch::Tensor generalised_geodesic3d_cpu(
-    torch::Tensor &image, 
-    const torch::Tensor &mask, 
-    const std::vector<float> &spacing, 
-    const float &v, 
-    const float &l_grad, 
-    const float &l_eucl, 
+    torch::Tensor &image,
+    const torch::Tensor &mask,
+    const std::vector<float> &spacing,
+    const float &v,
+    const float &l_grad,
+    const float &l_eucl,
     const int &iterations);
 
 torch::Tensor generalised_geodesic2d(
-    torch::Tensor &image, 
-    const torch::Tensor &mask, 
-    const float &v, 
-    const float &l_grad, 
-    const float &l_eucl, 
-    const int &iterations
-    );
+    torch::Tensor &image,
+    const torch::Tensor &mask,
+    const float &v,
+    const float &l_grad,
+    const float &l_eucl,
+    const int &iterations);
 
 torch::Tensor generalised_geodesic3d(
-    torch::Tensor &image, 
-    const torch::Tensor &mask, 
-    const std::vector<float> &spacing, 
-    const float &v, 
-    const float &l_grad, 
-    const float &l_eucl, 
-    const int &iterations
-    );
-
-
+    torch::Tensor &image,
+    const torch::Tensor &mask,
+    const std::vector<float> &spacing,
+    const float &v,
+    const float &l_grad,
+    const float &l_eucl,
+    const int &iterations);