WIP: DisplacementField class

src/deform_lib/registration/blocked_graph_cut_optimizer.inl

@@ -76,8 +76,7 @@ void BlockedGraphCutOptimizer<TUnaryTerm, TBinaryTerm, TSolver>
         if (_max_iteration_count != -1 && num_iterations >= _max_iteration_count)
             break;
 
-        // TODO: !!!!!
-        unary_fn.pre_iteration_hook(num_iterations, df.volume());
+        unary_fn.pre_iteration_hook(num_iterations, df);
 
         done = true;
 

src/deform_lib/registration/cost_functions/mutual_information.h

@@ -401,12 +401,12 @@ struct MIFunction : public SubFunction
      * entropy of the moving image and the joint entropy after each
      * iteration.
      */
-    virtual void pre_iteration_hook(const int iteration, const stk::VolumeFloat3& def)
+    virtual void pre_iteration_hook(const int iteration, const DisplacementField& df)
     {
         if (0 == iteration || 0 == _update_interval || iteration % _update_interval) {
             return;
         }
-        auto tmp = transform_volume(_moving, def, _interpolator);
+        auto tmp = transform_volume(_moving, df, _interpolator);
         _joint_entropy.update(_fixed, tmp);
         _entropy.update(tmp);
     }

src/deform_lib/registration/cost_functions/sub_function.h

@@ -6,6 +6,8 @@
 
 #include <limits>
 
+class DisplacementField;
+
 struct SubFunction
 {
     virtual ~SubFunction() {}
@@ -23,10 +25,10 @@ struct SubFunction
      * \param iteration Number of the iteration just completed.
      * \param def Deformation field at the end of the iteration.
      */
-    virtual void pre_iteration_hook(const int iteration, const stk::VolumeFloat3& def)
+    virtual void pre_iteration_hook(const int iteration, const DisplacementField& df)
     {
         (void) iteration;
-        (void) def;
+        (void) df;
     }
 
     /*!

src/deform_lib/registration/cost_functions/unary_function.h

@@ -41,10 +41,10 @@ struct UnaryFunction
         return mask_value * sum;
     }
 
-    void pre_iteration_hook(const int iteration, const stk::VolumeFloat3& def)
+    void pre_iteration_hook(const int iteration, const DisplacementField& df)
     {
         for (auto& fn : _functions) {
-            fn.function->pre_iteration_hook(iteration, def);
+            fn.function->pre_iteration_hook(iteration, df);
         }
     }
 

src/deform_lib/registration/displacement_field.h

@@ -9,22 +9,22 @@ class DisplacementField
 {
 public:
     DisplacementField(
-        Settings::UpdateRule update_rule,
-        const stk::VolumeFloat3& df
+        const stk::VolumeFloat3& df,
+        Settings::UpdateRule update_rule = Settings::UpdateRule_Additive
     ) :
         _update_rule(update_rule),
         _df(df)
     {
     }
     ~DisplacementField() {}
 
-    inline float3 get(const int3& p)
+    inline float3 get(const int3& p) const
     {
         return _df(p);
     }
 
     // delta : Delta in world space (mm)
-    inline float3 get(const int3& p, const float3& delta)
+    inline float3 get(const int3& p, const float3& delta) const
     {
         if (_update_rule == Settings::UpdateRule_Compositive) {
             // Convert delta to fixed image space
@@ -52,6 +52,13 @@ class DisplacementField
         _df(p) = get(p, delta);
     }
 
+    // p : Index in displacement field
+    // Returns coordinates in world space
+    inline float3 transform_index(const int3& p)
+    {
+        return _df.index2point(p) + get(p);
+    }
+
     dim3 size() const
     {
         return _df.size();
@@ -63,6 +70,24 @@ class DisplacementField
         return _df;
     }
 
+    inline const float3& operator()(int x, int y, int z) const
+    {
+        return _df(p);
+    }
+    inline float3& operator()(int x, int y, int z)
+    {
+        return _df(p);
+    }
+
+    inline const float3& operator()(const int3& p) const
+    {
+        return _df(p);
+    }
+    inline float3& operator()(const int3& p)
+    {
+        return _df(p);
+    }
+
 private:
     Settings::UpdateRule _update_rule;
     stk::VolumeFloat3 _df;

src/deform_lib/registration/registration_engine.cpp

@@ -618,7 +618,7 @@ stk::Volume RegistrationEngine::execute()
     #endif
 
     for (int l = _settings.num_pyramid_levels-1; l >= 0; --l) {
-        DisplacementField df(_settings.update_rule, _deformation_pyramid.volume(l));
+        DisplacementField df(_deformation_pyramid.volume(l), _settings.update_rule);
 
         std::vector<int3> neighborhood = determine_neighborhood(l);
 

src/deform_lib/registration/transform.cpp

@@ -1,3 +1,4 @@
+#include "displacement_field.h"
 #include "transform.h"
 
 #include <stk/common/log.h>
@@ -8,7 +9,7 @@ namespace
     template<typename TVoxelType>
     stk::VolumeHelper<TVoxelType> transform_volume_nn(
         const stk::VolumeHelper<TVoxelType>& src,
-        const stk::VolumeFloat3& def)
+        const DisplacementField& df)
     {
         // Transformed volume will have the same dimensions and properties (origin and spacing)
         //  as the deformation field and not the source image. The deformation field is inherited
@@ -18,28 +19,28 @@ namespace
         // There should be no requirements on src to have the same size, spacing and origin as
         //  the fixed image.
 
-        dim3 dims = def.size();
+        dim3 dims = df.size();
 
         stk::VolumeHelper<TVoxelType> out(dims);
-        out.copy_meta_from(def);
+        out.copy_meta_from(df);
 
         #define FAST_ROUND(x_) int(x_+0.5f)
 
         #pragma omp parallel for
         for (int z = 0; z < int(dims.z); ++z) {
-            for (int y = 0; y < int(dims.y); ++y) {
-                for (int x = 0; x < int(dims.x); ++x) {
-                    // [fixed] -> [world] -> [moving]
-                    const float3 moving_p = src.point2index(def.index2point(int3({x, y, z})) + def(x, y, z));
-
-                    out(x, y, z) = src.at(
-                        int(FAST_ROUND(moving_p.x)),
-                        int(FAST_ROUND(moving_p.y)),
-                        int(FAST_ROUND(moving_p.z)),
-                        stk::Border_Constant
-                    );
-                }
-            }
+        for (int y = 0; y < int(dims.y); ++y) {
+        for (int x = 0; x < int(dims.x); ++x) {
+            // [fixed] -> [world] -> [moving]
+            const float3 moving_p = src.point2index(df.transform_index(int3{x,y,z});
+
+            out(x, y, z) = src.at(
+                int(FAST_ROUND(moving_p.x)),
+                int(FAST_ROUND(moving_p.y)),
+                int(FAST_ROUND(moving_p.z)),
+                stk::Border_Constant
+            );
+        }
+        }
         }
         #undef FAST_ROUND
 
@@ -49,7 +50,7 @@ namespace
     template<typename TVoxelType>
     stk::VolumeHelper<TVoxelType> transform_volume_linear(
         const stk::VolumeHelper<TVoxelType>& src,
-        const stk::VolumeFloat3& def)
+        const DisplacementField& df)
     {
         // Transformed volume will have the same dimensions and properties (origin and spacing)
         //  as the deformation field and not the source image. The deformation field is inherited
@@ -77,28 +78,28 @@ namespace
     }
 }
 
-stk::Volume transform_volume(const stk::Volume& src, const stk::VolumeFloat3& def, transform::Interp interp)
+stk::Volume transform_volume(const stk::Volume& src, const DisplacementField& df, transform::Interp interp)
 {
     if (interp == transform::Interp_NN) {
         if (src.voxel_type() == stk::Type_Float) {
-            return transform_volume_nn<float>(src, def);
+            return transform_volume_nn<float>(src, df);
         }
         else if (src.voxel_type() == stk::Type_Double) {
-            return transform_volume_nn<double>(src, def);
+            return transform_volume_nn<double>(src, df);
         }
         else if (src.voxel_type() == stk::Type_UChar) {
-            return transform_volume_nn<uint8_t>(src, def);
+            return transform_volume_nn<uint8_t>(src, df);
         }
         else {
             LOG(Error) << "transform_volume: Unsupported volume type (type: " << src.voxel_type() << ")";
         }
     }
     else if (interp == transform::Interp_Linear) {
         if (src.voxel_type() == stk::Type_Float) {
-            return transform_volume_linear<float>(src, def);
+            return transform_volume_linear<float>(src, df);
         }
         else if (src.voxel_type() == stk::Type_Double) {
-            return transform_volume_linear<double>(src, def);
+            return transform_volume_linear<double>(src, df);
         }
         else {
             LOG(Error) << "transform_volume: Unsupported volume type (type: " << src.voxel_type() << ")";
@@ -109,3 +110,8 @@ stk::Volume transform_volume(const stk::Volume& src, const stk::VolumeFloat3& de
     }
     return stk::Volume();
 }
+
+stk::Volume transform_volume(const stk::Volume& src, const stk::VolumeFloat3& df, transform::Interp interp)
+{
+    return transform_volume(src, DisplacementField(df), interp);
+}

src/deform_lib/registration/transform.h

@@ -2,6 +2,8 @@
 
 #include <stk/image/volume.h>
 
+class DisplacementField;
+
 namespace transform
 {
     enum Interp : uint8_t
@@ -13,6 +15,12 @@ namespace transform
 
 stk::Volume transform_volume(
     const stk::Volume& src,
-    const stk::VolumeFloat3& def,
+    const DisplacementField& df,
+    transform::Interp i = transform::Interp_Linear
+);
+
+stk::Volume transform_volume(
+    const stk::Volume& src,
+    const stk::VolumeFloat3& df,
     transform::Interp i = transform::Interp_Linear
 );