Automatically adapt to image dimensionality

src/deform_lib/registration/blocked_graph_cut_optimizer.h

@@ -8,6 +8,7 @@
 #include "../config.h"
 #include "settings.h"
 
+#include <vector>
 
 template<
     typename TUnaryTerm,
@@ -18,6 +19,7 @@ class BlockedGraphCutOptimizer
 {
 public:
     BlockedGraphCutOptimizer(
+        const std::vector<int3>& neighborhood,
         const int3& block_size,
         double block_energy_epsilon,
         int max_iteration_count,
@@ -50,7 +52,7 @@ class BlockedGraphCutOptimizer
         stk::VolumeFloat3& def
     );
 
-    int3 _neighbors[6];
+    std::vector<int3> _neighborhood;
     int3 _block_size;
     double _block_energy_epsilon;
 

src/deform_lib/registration/blocked_graph_cut_optimizer.inl

@@ -13,21 +13,17 @@ template<
     typename TSolver
     >
 BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm, TSolver>::BlockedGraphCutOptimizer(
+    const std::vector<int3>& neighborhood,
     const int3& block_size,
     double block_energy_epsilon,
     int max_iteration_count,
     Settings::UpdateRule update_rule) :
+    _neighborhood(neighborhood),
     _block_size(block_size),
     _block_energy_epsilon(block_energy_epsilon),
     _max_iteration_count(max_iteration_count),
     _update_rule(update_rule)
 {
-    _neighbors[0] = {1, 0, 0};
-    _neighbors[1] = {-1, 0, 0};
-    _neighbors[2] = {0, 1, 0};
-    _neighbors[3] = {0, -1, 0};
-    _neighbors[4] = {0, 0, 1};
-    _neighbors[5] = {0, 0, -1};
 }
 template<
     typename TUnaryTerm,
@@ -126,9 +122,8 @@ void BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm, TSolver>::execute(
                     int3 block_p{block_x, block_y, block_z};
 
                     bool need_update = change_flags.is_block_set(block_p, use_shift == 1);
-                    int n_count = 6; // Neighbors
-                    for (int n = 0; n < n_count; ++n) {
-                        int3 neighbor = block_p + _neighbors[n];
+                    for (int3 n : _neighborhood) {
+                        int3 neighbor = block_p + n;
                         if (0 <= neighbor.x && neighbor.x < real_block_count.x &&
                             0 <= neighbor.y && neighbor.y < real_block_count.y &&
                             0 <= neighbor.z && neighbor.z < real_block_count.z) {
@@ -141,12 +136,13 @@ void BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm, TSolver>::execute(
                     }
 
                     bool block_changed = false;
-                    for (int n = 0; n < n_count; ++n) {
+
+                    for (int3 n : _neighborhood) {
                         // delta in [mm]
                         float3 delta {
-                            step_size.x * _neighbors[n].x,
-                            step_size.y * _neighbors[n].y,
-                            step_size.z * _neighbors[n].z
+                            step_size.x * n.x,
+                            step_size.y * n.y,
+                            step_size.z * n.z
                         };
 
                         block_changed |= do_block(

src/deform_lib/registration/gpu_registration_engine.cpp

@@ -277,18 +277,16 @@ stk::Volume GpuRegistrationEngine::execute()
 
         stk::GpuVolume base = _fixed_pyramids[0].volume(0);
 
-        // TODO: Until we get a proper GpuVolume::fill()
         stk::VolumeFloat4 initial(base.size(), float4{0, 0, 0, 0});
-        initial.set_origin(base.origin());
-        initial.set_spacing(base.spacing());
-        initial.set_direction(base.direction());
-
+        initial.copy_meta_from(base);
         set_initial_deformation(initial);
     }
 
     for (int l = _settings.num_pyramid_levels-1; l >= 0; --l) {
         stk::GpuVolume df = _deformation_pyramid.volume(l);
 
+        std::vector<int3> neighborhood = determine_neighborhood(l);
+
         if (l >= _settings.pyramid_stop_level) {
             LOG(Info) << "Performing registration level " << l;
 
@@ -300,6 +298,7 @@ stk::Volume GpuRegistrationEngine::execute()
 
             if (_settings.solver == Settings::Solver_ICM) {
                 HybridGraphCutOptimizer<ICMSolver<double>> optimizer(
+                    neighborhood,
                     _settings.levels[l],
                     _settings.update_rule,
                     unary_fn,
@@ -313,6 +312,7 @@ stk::Volume GpuRegistrationEngine::execute()
 #if defined(DF_ENABLE_GCO)
             else if (_settings.solver == Settings::Solver_GCO) {
                 HybridGraphCutOptimizer<GCOSolver<double>> optimizer(
+                    neighborhood,
                     _settings.levels[l],
                     _settings.update_rule,
                     unary_fn,
@@ -327,6 +327,7 @@ stk::Volume GpuRegistrationEngine::execute()
 #if defined(DF_ENABLE_GRIDCUT)
             else if (_settings.solver == Settings::Solver_GridCut) {
                 HybridGraphCutOptimizer<GridCutSolver<double>> optimizer(
+                    neighborhood,
                     _settings.levels[l],
                     _settings.update_rule,
                     unary_fn,
@@ -339,7 +340,6 @@ stk::Volume GpuRegistrationEngine::execute()
             }
 #endif
 
-
         }
         else {
             LOG(Info) << "Skipping level " << l;
@@ -359,3 +359,51 @@ stk::Volume GpuRegistrationEngine::execute()
 
     return volume_float4_to_float3(_deformation_pyramid.volume(0).download());
 }
+std::vector<int3> GpuRegistrationEngine::determine_neighborhood(int level) const
+{
+    // Identical to RegistrationEngine::determine_neighborhood with the exception
+    //  of working on GpuVolumePyramid
+
+    std::vector<int3> neighborhood;
+
+    dim3 dim_size {0, 0, 0};
+
+    for (int i = 0; i < (int) _fixed_pyramids.size(); ++i) {
+        stk::Volume fixed;
+        stk::Volume moving;
+
+        if (_fixed_pyramids[i].levels() > 0)
+            fixed = _fixed_pyramids[i].volume(level);
+        if (_moving_pyramids[i].levels() > 0)
+            moving = _moving_pyramids[i].volume(level);
+
+        dim_size = {
+            std::max(dim_size.x, fixed.size().x),
+            std::max(dim_size.y, fixed.size().y),
+            std::max(dim_size.z, fixed.size().z)
+        };
+
+        dim_size = {
+            std::max(dim_size.x, moving.size().x),
+            std::max(dim_size.y, moving.size().y),
+            std::max(dim_size.z, moving.size().z)
+        };
+    }
+
+    if (dim_size.x > 1) {
+        neighborhood.push_back({1, 0, 0});
+        neighborhood.push_back({-1, 0, 0});
+    }
+
+    if (dim_size.y > 1) {
+        neighborhood.push_back({0, 1, 0});
+        neighborhood.push_back({0, -1, 0});
+    }
+
+    if (dim_size.z > 1) {
+        neighborhood.push_back({0, 0, 1});
+        neighborhood.push_back({0, 0, -1});
+    }
+
+    return neighborhood;
+}

src/deform_lib/registration/gpu_registration_engine.h

@@ -46,6 +46,11 @@ class GpuRegistrationEngine
     stk::Volume execute();
 
 private:
+    /// Determines the neighborhood to apply for the given pyramid level.
+    /// The neighborhood is a set of unit vectors in which directions
+    /// to apply the optimization. E.g. ((-1,0,0), (1,0,0), (0,-1,0), ...)
+    std::vector<int3> determine_neighborhood(int level) const;
+
     /// Builds a unary function for the specified pyramid level
     void build_unary_function(int level, GpuUnaryFunction& unary_fn);
 

src/deform_lib/registration/hybrid_graph_cut_optimizer.h

@@ -10,6 +10,7 @@
 #include <atomic>
 #include <deque>
 #include <mutex>
+#include <vector>
 
 class GpuBinaryFunction;
 class GpuUnaryFunction;
@@ -23,6 +24,7 @@ class HybridGraphCutOptimizer
 {
 public:
     HybridGraphCutOptimizer(
+        const std::vector<int3>& neighborhood,
         const Settings::Level& settings,
         Settings::UpdateRule update_rule,
         GpuUnaryFunction& unary_fn,
@@ -86,6 +88,8 @@ class HybridGraphCutOptimizer
     // Note: This resets the cost buffers
     double calculate_energy();
 
+    std::vector<int3> _neighborhood;
+
     const Settings::Level& _settings;
     Settings::UpdateRule _update_rule;
 

src/deform_lib/registration/hybrid_graph_cut_optimizer.inl

@@ -10,27 +10,18 @@
 
 #include <iomanip>
 
-namespace {
-    const int _neighbor_count = 6;
-    int3 _neighbors[] = {
-        {1, 0, 0},
-        {-1, 0, 0},
-        {0, 1, 0},
-        {0, -1, 0},
-        {0, 0, 1},
-        {0, 0, -1}
-    };
-}
 
 template<typename TSolver>
 HybridGraphCutOptimizer<TSolver>::HybridGraphCutOptimizer(
+    const std::vector<int3>& neighborhood,
     const Settings::Level& settings,
     Settings::UpdateRule update_rule,
     GpuUnaryFunction& unary_fn,
     GpuBinaryFunction& binary_fn,
     stk::GpuVolume& df,
     WorkerPool& worker_pool,
     std::vector<stk::cuda::Stream>& stream_pool) :
+    _neighborhood(neighborhood),
     _settings(settings),
     _update_rule(update_rule),
     _worker_pool(worker_pool),
@@ -125,14 +116,14 @@ void HybridGraphCutOptimizer<TSolver>::execute()
         for (int black_or_red = 0; black_or_red < (num_blocks > 1 ? 2 : 1); black_or_red++) {
         PROFILER_SCOPE("red_black", 0xFF339955);
 
-        for (int n = 0; n < _neighbor_count; ++n) {
+        for (int3 n : _neighborhood) {
         PROFILER_SCOPE("step", 0xFFAA6FE2);
 
         // delta in [mm]
         _current_delta = {
-            _settings.step_size.x * _neighbors[n].x,
-            _settings.step_size.y * _neighbors[n].y,
-            _settings.step_size.z * _neighbors[n].z
+            _settings.step_size.x * n.x,
+            _settings.step_size.y * n.y,
+            _settings.step_size.z * n.z
         };
 
         // Queue all blocks
@@ -153,8 +144,8 @@ void HybridGraphCutOptimizer<TSolver>::execute()
 
             bool need_update = _block_change_flags.is_block_set(block_idx,
                                                                 use_shift == 1);
-            for (int i = 0; i < _neighbor_count; ++i) {
-                int3 neighbor = block_idx + _neighbors[i];
+            for (int3 n2 : _neighborhood) {
+                int3 neighbor = block_idx + n2;
                 if (0 <= neighbor.x && neighbor.x < real_block_count.x &&
                     0 <= neighbor.y && neighbor.y < real_block_count.y &&
                     0 <= neighbor.z && neighbor.z < real_block_count.z) {

src/deform_lib/registration/registration_engine.cpp

@@ -618,6 +618,8 @@ stk::Volume RegistrationEngine::execute()
     for (int l = _settings.num_pyramid_levels-1; l >= 0; --l) {
         stk::VolumeFloat3 def = _deformation_pyramid.volume(l);
 
+        std::vector<int3> neighborhood = determine_neighborhood(l);
+
         if (l >= _settings.pyramid_stop_level) {
             LOG(Info) << "Performing registration level " << l;
 
@@ -640,6 +642,7 @@ stk::Volume RegistrationEngine::execute()
             if (_settings.solver == Settings::Solver_ICM) {
                 BlockedGraphCutOptimizer<UnaryFunction, Regularizer, ICMSolver<double>>
                     optimizer(
+                        neighborhood,
                         _settings.levels[l].block_size,
                         _settings.levels[l].block_energy_epsilon,
                         _settings.levels[l].max_iteration_count,
@@ -651,6 +654,7 @@ stk::Volume RegistrationEngine::execute()
             else if (_settings.solver == Settings::Solver_GCO) {
                 BlockedGraphCutOptimizer<UnaryFunction, Regularizer, GCOSolver<double>>
                     optimizer(
+                        neighborhood,
                         _settings.levels[l].block_size,
                         _settings.levels[l].block_energy_epsilon,
                         _settings.levels[l].max_iteration_count,
@@ -663,6 +667,7 @@ stk::Volume RegistrationEngine::execute()
             else if (_settings.solver == Settings::Solver_GridCut) {
                 BlockedGraphCutOptimizer<UnaryFunction, Regularizer, GridCutSolver<double>>
                     optimizer(
+                        neighborhood,
                         _settings.levels[l].block_size,
                         _settings.levels[l].block_energy_epsilon,
                         _settings.levels[l].max_iteration_count,
@@ -702,6 +707,63 @@ stk::Volume RegistrationEngine::deformation_field(int level)
 {
     return _deformation_pyramid.volume(level);
 }
+std::vector<int3> RegistrationEngine::determine_neighborhood(int level) const
+{
+    /*
+        To enable registration of images with lower dimensionality than 3 we
+        automatically adopt the neighborhood for the optimization. Given a set
+        of images (all fixed and moving images) we determine the dimensionality
+        for the registration by assuming that if all images are of size 1 in a
+        specific dimension, that dimension is not in play.
+
+        E.g., given an input of only WxHx1 images, we trigger 2D registration
+        by setting the search neighborhood to only the X and Y axes.
+    */
+
+    std::vector<int3> neighborhood;
+
+    dim3 dim_size {0, 0, 0};
+
+    for (int i = 0; i < (int) _fixed_pyramids.size(); ++i) {
+        stk::Volume fixed;
+        stk::Volume moving;
+
+        if (_fixed_pyramids[i].levels() > 0)
+            fixed = _fixed_pyramids[i].volume(level);
+        if (_moving_pyramids[i].levels() > 0)
+            moving = _moving_pyramids[i].volume(level);
+
+        dim_size = {
+            std::max(dim_size.x, fixed.size().x),
+            std::max(dim_size.y, fixed.size().y),
+            std::max(dim_size.z, fixed.size().z)
+        };
+
+        dim_size = {
+            std::max(dim_size.x, moving.size().x),
+            std::max(dim_size.y, moving.size().y),
+            std::max(dim_size.z, moving.size().z)
+        };
+    }
+
+    if (dim_size.x > 1) {
+        neighborhood.push_back({1, 0, 0});
+        neighborhood.push_back({-1, 0, 0});
+    }
+
+    if (dim_size.y > 1) {
+        neighborhood.push_back({0, 1, 0});
+        neighborhood.push_back({0, -1, 0});
+    }
+
+    if (dim_size.z > 1) {
+        neighborhood.push_back({0, 0, 1});
+        neighborhood.push_back({0, 0, -1});
+    }
+
+    return neighborhood;
+}
+
 #ifdef DF_OUTPUT_DEBUG_VOLUMES
 void RegistrationEngine::upsample_and_save(int level)
 {

src/deform_lib/registration/registration_engine.h

@@ -52,6 +52,11 @@ class RegistrationEngine
     stk::Volume deformation_field(int level);
 
 private:
+    /// Determines the neighborhood to apply for the given pyramid level.
+    /// The neighborhood is a set of unit vectors in which directions
+    /// to apply the optimization. E.g. ((-1,0,0), (1,0,0), (0,-1,0), ...)
+    std::vector<int3> determine_neighborhood(int level) const;
+
 #ifdef DF_OUTPUT_DEBUG_VOLUMES
     /// Upsamples the deformation field at the given level and saves it
     void upsample_and_save(int level);