Support locally weighting direction field losses

apps/src/vc_grow_seg_from_seed.cpp

@@ -100,8 +100,18 @@ static auto load_direction_fields(json const&params, ChunkCache *chunk_cache, st
                 direction_dss.push_back(z5::filesystem::openDataset(dirs_ds_handle));
             }
             std::cout << "direction field dataset shape " << direction_dss.front()->shape() << std::endl;
+            std::unique_ptr<z5::Dataset> maybe_weight_ds;
+            if (direction_field.contains("weight_zarr")) {
+                std::string const weight_zarr_path = direction_field["weight_zarr"];
+                z5::filesystem::handle::Group weight_group(weight_zarr_path);
+                z5::filesystem::handle::Dataset weight_ds_handle(weight_group, std::to_string(ome_scale), ".");
+                maybe_weight_ds = z5::filesystem::openDataset(weight_ds_handle);
+            }
             std::string const unique_id = std::to_string(std::hash<std::string>{}(dirs_group.path().string() + std::to_string(ome_scale)));
-            direction_fields.emplace_back(direction, std::make_unique<Chunked3dVec3fFromUint8>(std::move(direction_dss), scale_factor, chunk_cache, cache_root, unique_id));
+            direction_fields.emplace_back(
+                direction,
+                std::make_unique<Chunked3dVec3fFromUint8>(std::move(direction_dss), scale_factor, chunk_cache, cache_root, unique_id),
+                maybe_weight_ds ? std::make_unique<Chunked3dFloatFromUint8>(std::move(maybe_weight_ds), scale_factor, chunk_cache, cache_root, unique_id + "_conf") : std::unique_ptr<Chunked3dFloatFromUint8>());
         }
     }
     return direction_fields;

core/include/vc/core/types/ChunkedTensor.hpp

@@ -623,6 +623,36 @@ class CachedChunked3dInterpolator
     std::vector<int>    _shape;
 };
 
+struct Chunked3dFloatFromUint8
+{
+    Chunked3dFloatFromUint8(std::unique_ptr<z5::Dataset> &&ds, float scale, ChunkCache *cache, std::string const &cache_root, std::string const &unique_id) :
+        _passthrough{unique_id},
+        _x(_passthrough, ds.get(), cache, cache_root),
+        _scale(scale),  // multiplying by this maps indices of the 'canonical' volume to indices of our dataset
+        _ds(std::move(ds))  // take ownership of the dataset, as Chunked3d accesses it through a bare pointer
+    {
+    }
+
+    float operator()(cv::Vec3d p)
+    {
+        // p has zyx ordering!
+        p *= _scale;
+        cv::Vec3i i{lround(p[0]), lround(p[1]), lround(p[2])};
+        uint8_t x = _x(i) ;
+        return float{x} / 255.f;
+    }
+
+    float operator()(double z, double y, double x)
+    {
+        return operator()({z,y,x});
+    }
+
+    passTroughComputor _passthrough;
+    Chunked3d<uint8_t, passTroughComputor> _x;
+    float _scale;
+    std::unique_ptr<z5::Dataset> _ds;
+};
+
 struct Chunked3dVec3fFromUint8
 {
     Chunked3dVec3fFromUint8(std::vector<std::unique_ptr<z5::Dataset>> &&dss, float scale, ChunkCache *cache, std::string const &cache_root, std::string const &unique_id) :

core/include/vc/core/util/CostFunctions.hpp

@@ -393,7 +393,8 @@ struct SpaceLineLossAcc {
 };
 
 struct FiberDirectionLoss {
-    FiberDirectionLoss(Chunked3dVec3fFromUint8 &fiber_dirs, float w) : _fiber_dirs(fiber_dirs), _w(w) {};
+    FiberDirectionLoss(Chunked3dVec3fFromUint8 &fiber_dirs, Chunked3dFloatFromUint8 *maybe_weights, float w) :
+        _fiber_dirs(fiber_dirs), _maybe_weights(maybe_weights), _w(w) {};
     template <typename E>
     bool operator()(const E* const l_base, const E* const l_u_off, E* residual) const {
 
@@ -417,7 +418,9 @@ struct FiberDirectionLoss {
         E const patch_u_dist = sqrt(patch_u_disp_zyx[0] * patch_u_disp_zyx[0] + patch_u_disp_zyx[1] * patch_u_disp_zyx[1] + patch_u_disp_zyx[2] * patch_u_disp_zyx[2]);
         E const abs_dot = abs(patch_u_disp_zyx[0] * fiber_dir_zyx[0] + patch_u_disp_zyx[1] * fiber_dir_zyx[1] + patch_u_disp_zyx[2] * fiber_dir_zyx[2]) / patch_u_dist;
 
-        residual[0] = E(_w) * (E(1) - abs_dot);
+        E const weight_at_point = _maybe_weights ? E((*_maybe_weights)(unjet(l_base[2]), unjet(l_base[1]), unjet(l_base[0]))) : E(1);
+
+        residual[0] = E(_w) * (E(1) - abs_dot) * weight_at_point;
 
         return true;
     }
@@ -427,15 +430,17 @@ struct FiberDirectionLoss {
 
     float _w;
     Chunked3dVec3fFromUint8 &_fiber_dirs;
+    Chunked3dFloatFromUint8 *_maybe_weights;
 
-    static ceres::CostFunction* Create(Chunked3dVec3fFromUint8 &fiber_dirs, float w = 1.0)
+    static ceres::CostFunction* Create(Chunked3dVec3fFromUint8 &fiber_dirs, Chunked3dFloatFromUint8 *maybe_weights, float w = 1.0)
     {
-        return new ceres::AutoDiffCostFunction<FiberDirectionLoss, 1, 3, 3>(new FiberDirectionLoss(fiber_dirs, w));
+        return new ceres::AutoDiffCostFunction<FiberDirectionLoss, 1, 3, 3>(new FiberDirectionLoss(fiber_dirs, maybe_weights, w));
     }
 };
 
 struct NormalDirectionLoss {
-    NormalDirectionLoss(Chunked3dVec3fFromUint8 &normal_dirs, float w) : _normal_dirs(normal_dirs), _w(w) {};
+    NormalDirectionLoss(Chunked3dVec3fFromUint8 &normal_dirs, Chunked3dFloatFromUint8 *maybe_weights, float w) :
+        _normal_dirs(normal_dirs), _maybe_weights(maybe_weights), _w(w) {};
     template <typename E>
     bool operator()(const E* const l_base, const E* const l_u_off, const E* const l_v_off, E* residual) const {
 
@@ -472,7 +477,9 @@ struct NormalDirectionLoss {
 
         E const abs_dot = abs(patch_normal_zyx[0] * target_normal_zyx[0] + patch_normal_zyx[1] * target_normal_zyx[1] + patch_normal_zyx[2] * target_normal_zyx[2]) / patch_normal_length;
 
-        residual[0] = E(_w) * (E(1) - abs_dot);
+        E const weight_at_point = _maybe_weights ? E((*_maybe_weights)(unjet(l_base[2]), unjet(l_base[1]), unjet(l_base[0]))) : E(1);
+
+        residual[0] = E(_w) * (E(1) - abs_dot) * weight_at_point;
 
         return true;
     }
@@ -482,9 +489,10 @@ struct NormalDirectionLoss {
 
     float _w;
     Chunked3dVec3fFromUint8 &_normal_dirs;
+    Chunked3dFloatFromUint8 *_maybe_weights;
 
-    static ceres::CostFunction* Create(Chunked3dVec3fFromUint8 &normal_dirs, float w = 1.0)
+    static ceres::CostFunction* Create(Chunked3dVec3fFromUint8 &normal_dirs, Chunked3dFloatFromUint8 *maybe_weights, float w = 1.0)
     {
-        return new ceres::AutoDiffCostFunction<NormalDirectionLoss, 1, 3, 3, 3>(new NormalDirectionLoss(normal_dirs, w));
+        return new ceres::AutoDiffCostFunction<NormalDirectionLoss, 1, 3, 3, 3>(new NormalDirectionLoss(normal_dirs, maybe_weights, w));
     }
 };

core/include/vc/tracer/Tracer.hpp

@@ -1,12 +1,14 @@
 #pragma once
 #include "vc/core/util/Surface.hpp"
 
+struct Chunked3dFloatFromUint8;
 struct Chunked3dVec3fFromUint8;
 
 struct DirectionField
 {
     std::string direction;
     std::unique_ptr<Chunked3dVec3fFromUint8> field_ptr;
+    std::unique_ptr<Chunked3dFloatFromUint8> weight_ptr;
 };
 
 QuadSurface *grow_surf_from_surfs(SurfaceMeta *seed, const std::vector<SurfaceMeta*> &surfs_v, const nlohmann::json &params, float voxelsize = 1.0);

core/src/GrowPatch.cpp

@@ -225,15 +225,15 @@ int gen_direction_loss(ceres::Problem &problem, const cv::Vec2i &p, const int of
         if (field.direction == "horizontal") {
             if (!loc_valid(state(p_off_horz)))
                 continue;
-            problem.AddResidualBlock(FiberDirectionLoss::Create(*field.field_ptr, w), nullptr, &loc(p)[0], &loc(p_off_horz)[0]);
+            problem.AddResidualBlock(FiberDirectionLoss::Create(*field.field_ptr, field.weight_ptr.get(), w), nullptr, &loc(p)[0], &loc(p_off_horz)[0]);
         } else if (field.direction == "vertical") {
             if (!loc_valid(state(p_off_vert)))
                 continue;
-            problem.AddResidualBlock(FiberDirectionLoss::Create(*field.field_ptr, w), nullptr, &loc(p)[0], &loc(p_off_vert)[0]);
+            problem.AddResidualBlock(FiberDirectionLoss::Create(*field.field_ptr, field.weight_ptr.get(), w), nullptr, &loc(p)[0], &loc(p_off_vert)[0]);
         } else if (field.direction == "normal") {
             if (!loc_valid(state(p_off_horz)) || !loc_valid(state(p_off_vert)))
                 continue;
-            problem.AddResidualBlock(NormalDirectionLoss::Create(*field.field_ptr, w), nullptr, &loc(p)[0], &loc(p_off_horz)[0], &loc(p_off_vert)[0]);
+            problem.AddResidualBlock(NormalDirectionLoss::Create(*field.field_ptr, field.weight_ptr.get(), w), nullptr, &loc(p)[0], &loc(p_off_horz)[0], &loc(p_off_vert)[0]);
         } else {
             assert(false);
         }