Merge pull request #103 from simeks/gridcut

GridCut

CMakeLists.txt

@@ -47,6 +47,17 @@ option(DF_NON_PORTABLE "Enable non-portable optimisations" ON)
 option(DF_ENABLE_MICROPROFILE "Enable microprofile profiler" OFF)
 option(DF_ENABLE_NVTOOLSEXT "Enable nvtoolsext profiler" OFF)
 
+option(DF_ENABLE_GCO "Enable GCO solver" ON)
+option(DF_ENABLE_GRIDCUT "Enable GridCut solver" OFF)
+
+if (DF_ENABLE_GCO)
+    add_definitions(-DDF_ENABLE_GCO)
+endif()
+
+if (DF_ENABLE_GRIDCUT)
+    add_definitions(-DDF_ENABLE_GRIDCUT)
+endif()
+
 if (DF_ENABLE_MICROPROFILE AND DF_ENABLE_NVTOOLSEXT)
     message(FATAL_ERROR "Can't enable multiple profilers at the same time")
 endif()

README.md

@@ -105,6 +105,7 @@ To perform a registration using the standalone executable
 pyramid_levels: 6
 pyramid_stop_level: 0
 regularize_initial_displacement: false
+solver: gco
 constraints_weight: 1000.0
 landmarks_weight: 1.0
 landmarks_decay: 2.0
@@ -162,6 +163,10 @@ specifies that the registration should not be run on the original resolution
 displacement in the regularization. This is disabled by default, meaning the
 initial displacement have no effect on the regularization.
 
+`solver` selects which solver to use for the energy minimization. Available
+solvers are `gco`, `gridcut`, and `icm`. Note: deform needs to be compiled with
+`DF_ENABLE_GCO` and `DF_ENABLE_GRIDCUT` for `gco` and `gridcut` to be enabled.
+
 `constraints_weight` sets the weight that is applied for constrained voxels. A
 really high value means hard constraints while a lower value may allow
 constraints to move a certain amount. Cost for constrained voxels are applied

src/deform_lib/CMakeLists.txt

@@ -16,12 +16,6 @@ set(DEFORM_LIB_SRCS
     "filters/resample.cpp"
     "filters/resample.h"
 
-    "graph_cut/graph_cut.h"
-    "graph_cut/graph_cut.inl"
-
-    "graph_cut/graph_cut.h"
-    "graph_cut/graph_cut.inl"
-
     "registration/block_change_flags.cpp"
     "registration/block_change_flags.h"
     "registration/blocked_graph_cut_optimizer.h"
@@ -51,6 +45,21 @@ set(DEFORM_LIB_SRCS
     "version.h"
 )
 
+if (DF_ENABLE_GCO)
+    set(DEFORM_LIB_SRCS
+        ${DEFORM_LIB_SRCS}
+        "solver/gco_solver.h"
+        "solver/gco_solver.inl"
+    )
+endif ()
+
+if (DF_ENABLE_GRIDCUT)
+    set(DEFORM_LIB_SRCS
+        ${DEFORM_LIB_SRCS}
+        "solver/gridcut_solver.h"
+    )
+endif()
+
 if (DF_USE_CUDA)
     set(DEFORM_LIB_SRCS
         ${DEFORM_LIB_SRCS}
@@ -77,7 +86,7 @@ if (DF_USE_CUDA)
         "registration/gpu_registration_engine.h"
         "registration/gpu_volume_pyramid.cpp"
         "registration/gpu_volume_pyramid.h"
-        "registration/hybrid_graph_cut_optimizer.cpp"
+        "registration/hybrid_graph_cut_optimizer.inl"
         "registration/hybrid_graph_cut_optimizer.cu"
         "registration/hybrid_graph_cut_optimizer.h"
     )
@@ -99,7 +108,15 @@ add_dependencies(deform_lib FetchGitVersion)
 
 target_include_directories(deform_lib PUBLIC ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES} "${CMAKE_BINARY_DIR}/src")
 
-target_link_libraries(deform_lib stk yaml-cpp gco)
+target_link_libraries(deform_lib stk yaml-cpp)
+
+if (DF_ENABLE_GCO)
+    target_link_libraries(deform_lib gco)
+endif()
+
+if (DF_ENABLE_GRIDCUT)
+    target_link_libraries(deform_lib gridcut)
+endif()
 
 if (DF_ENABLE_NVTOOLSEXT)
     target_link_libraries(deform_lib ${NvToolsExt_LIBRARIES})

src/deform_lib/registration/blocked_graph_cut_optimizer.h

@@ -10,7 +10,8 @@
 
 template<
     typename TUnaryTerm,
-    typename TBinaryTerm
+    typename TBinaryTerm,
+    typename TSolver
     >
 class BlockedGraphCutOptimizer
 {

src/deform_lib/registration/blocked_graph_cut_optimizer.inl

@@ -1,6 +1,5 @@
 #include "block_change_flags.h"
 #include "../config.h"
-#include "../graph_cut/graph_cut.h"
 #include "../profiler/profiler.h"
 
 #include <stk/common/log.h>
@@ -9,9 +8,10 @@
 
 template<
     typename TUnaryTerm,
-    typename TBinaryTerm
->
-BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm>::BlockedGraphCutOptimizer(
+    typename TBinaryTerm,
+    typename TSolver
+    >
+BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm, TSolver>::BlockedGraphCutOptimizer(
     const int3& block_size,
     double block_energy_epsilon,
     int max_iteration_count) :
@@ -28,16 +28,18 @@ BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm>::BlockedGraphCutOptimizer(
 }
 template<
     typename TUnaryTerm,
-    typename TBinaryTerm
+    typename TBinaryTerm,
+    typename TSolver
     >
-BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm>::~BlockedGraphCutOptimizer()
+BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm, TSolver>::~BlockedGraphCutOptimizer()
 {
 }
 template<
     typename TUnaryTerm,
-    typename TBinaryTerm
+    typename TBinaryTerm,
+    typename TSolver
     >
-void BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm>::execute(
+void BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm, TSolver>::execute(
     TUnaryTerm& unary_fn,
     TBinaryTerm& binary_fn,
     float3 step_size,
@@ -182,9 +184,10 @@ void BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm>::execute(
 
 template<
     typename TUnaryTerm,
-    typename TBinaryTerm
+    typename TBinaryTerm,
+    typename TSolver
 >
-bool BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm>::do_block(
+bool BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm, TSolver>::do_block(
     TUnaryTerm& unary_fn,
     TBinaryTerm& binary_fn,
     const int3& block_p,
@@ -194,11 +197,11 @@ bool BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm>::do_block(
     stk::VolumeFloat3& def
 )
 {
-    dim3 dims = def.size();
+    using FlowType = typename TSolver::FlowType;
 
-    typedef double FlowType;
+    dim3 dims = def.size();
 
-    GraphCut<FlowType> graph(block_dims);
+    TSolver solver(block_dims);
 
     FlowType current_energy = 0;
     {
@@ -216,7 +219,7 @@ bool BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm>::do_block(
                     if (gx < 0 || gx >= int(dims.x) ||
                         gy < 0 || gy >= int(dims.y) ||
                         gz < 0 || gz >= int(dims.z)) {
-                        graph.add_term1(sub_x, sub_y, sub_z, 0, 0);
+                        solver.add_term1(sub_x, sub_y, sub_z, 0, 0);
                         continue;
                     }
 
@@ -267,7 +270,7 @@ bool BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm>::do_block(
                         f1 += binary_fn(p, def1 + delta, def2, step);
                     }
 
-                    graph.add_term1(sub_x, sub_y, sub_z, f0, f1);
+                    solver.add_term1(sub_x, sub_y, sub_z, f0, f1);
 
                     current_energy += f0;
 
@@ -278,7 +281,7 @@ bool BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm>::do_block(
                         double f01 = binary_fn(p, def1, def2 + delta, step);
                         double f10 = binary_fn(p, def1 + delta, def2, step);
 
-                        graph.add_term2(
+                        solver.add_term2(
                             sub_x, sub_y, sub_z,
                             sub_x + 1, sub_y, sub_z,
                             f_same, f01, f10, f_same);
@@ -292,7 +295,7 @@ bool BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm>::do_block(
                         double f01 = binary_fn(p, def1, def2 + delta, step);
                         double f10 = binary_fn(p, def1 + delta, def2, step);
 
-                        graph.add_term2(
+                        solver.add_term2(
                             sub_x, sub_y, sub_z,
                             sub_x, sub_y + 1, sub_z,
                             f_same, f01, f10, f_same);
@@ -306,7 +309,7 @@ bool BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm>::do_block(
                         double f01 = binary_fn(p, def1, def2 + delta, step);
                         double f10 = binary_fn(p, def1 + delta, def2, step);
 
-                        graph.add_term2(
+                        solver.add_term2(
                             sub_x, sub_y, sub_z,
                             sub_x, sub_y, sub_z + 1,
                             f_same, f01, f10, f_same);
@@ -318,11 +321,10 @@ bool BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm>::do_block(
         }
     }
 
-
     FlowType current_emin;
     {
         PROFILER_SCOPE("minimize", 0xFF985423);
-        current_emin = graph.minimize();
+        current_emin = solver.minimize();
     }
 
     bool changed_flag = false;
@@ -346,7 +348,7 @@ bool BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm>::do_block(
                         continue;
                     }
 
-                    if (graph.get_var(sub_x, sub_y, sub_z) == 1)
+                    if (solver.get_var(sub_x, sub_y, sub_z) == 1)
                     {
                         def(gx, gy, gz) = def(gx, gy, gz) + delta;
                         changed_flag = true;
@@ -362,9 +364,10 @@ bool BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm>::do_block(
 
 template<
     typename TUnaryTerm,
-    typename TBinaryTerm
+    typename TBinaryTerm,
+    typename TSolver
 >
-double BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm>::calculate_energy(
+double BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm, TSolver>::calculate_energy(
     TUnaryTerm& unary_fn,
     TBinaryTerm& binary_fn,
     stk::VolumeFloat3& def