From 38f28d9766ffae5c8d6668063111db835960d1f0 Mon Sep 17 00:00:00 2001 From: Paul Henderson Date: Thu, 4 Sep 2025 21:30:50 +0100 Subject: [PATCH 1/2] Support locally weighting direction field losses by values of another zarr --- apps/src/vc_grow_seg_from_seed.cpp | 12 +++++++- core/include/vc/core/types/ChunkedTensor.hpp | 30 ++++++++++++++++++++ core/include/vc/core/util/CostFunctions.hpp | 24 ++++++++++------ core/include/vc/tracer/Tracer.hpp | 2 ++ core/src/GrowPatch.cpp | 6 ++-- 5 files changed, 62 insertions(+), 12 deletions(-) diff --git a/apps/src/vc_grow_seg_from_seed.cpp b/apps/src/vc_grow_seg_from_seed.cpp index 84fc1d0b9..63e77f168 100644 --- a/apps/src/vc_grow_seg_from_seed.cpp +++ b/apps/src/vc_grow_seg_from_seed.cpp @@ -100,8 +100,18 @@ static auto load_direction_fields(json const¶ms, 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 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{}(dirs_group.path().string() + std::to_string(ome_scale))); - direction_fields.emplace_back(direction, std::make_unique(std::move(direction_dss), scale_factor, chunk_cache, cache_root, unique_id)); + direction_fields.emplace_back( + direction, + std::make_unique(std::move(direction_dss), scale_factor, chunk_cache, cache_root, unique_id), + maybe_weight_ds ? std::make_unique(std::move(maybe_weight_ds), scale_factor, chunk_cache, cache_root, unique_id + "_conf") : std::unique_ptr()); } } return direction_fields; diff --git a/core/include/vc/core/types/ChunkedTensor.hpp b/core/include/vc/core/types/ChunkedTensor.hpp index fb513324f..c40a2d158 100644 --- a/core/include/vc/core/types/ChunkedTensor.hpp +++ b/core/include/vc/core/types/ChunkedTensor.hpp @@ -623,6 +623,36 @@ class CachedChunked3dInterpolator std::vector _shape; }; +struct Chunked3dFloatFromUint8 +{ + Chunked3dFloatFromUint8(std::unique_ptr &&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} - 128.f) / 127.f; + } + + float operator()(double z, double y, double x) + { + return operator()({z,y,x}); + } + + passTroughComputor _passthrough; + Chunked3d _x; + float _scale; + std::unique_ptr _ds; +}; + struct Chunked3dVec3fFromUint8 { Chunked3dVec3fFromUint8(std::vector> &&dss, float scale, ChunkCache *cache, std::string const &cache_root, std::string const &unique_id) : diff --git a/core/include/vc/core/util/CostFunctions.hpp b/core/include/vc/core/util/CostFunctions.hpp index 80d5e7757..6eb2036d3 100644 --- a/core/include/vc/core/util/CostFunctions.hpp +++ b/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 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(new FiberDirectionLoss(fiber_dirs, w)); + return new ceres::AutoDiffCostFunction(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 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(new NormalDirectionLoss(normal_dirs, w)); + return new ceres::AutoDiffCostFunction(new NormalDirectionLoss(normal_dirs, maybe_weights, w)); } }; diff --git a/core/include/vc/tracer/Tracer.hpp b/core/include/vc/tracer/Tracer.hpp index 5dccdc9e5..06cd17d38 100644 --- a/core/include/vc/tracer/Tracer.hpp +++ b/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 field_ptr; + std::unique_ptr weight_ptr; }; QuadSurface *grow_surf_from_surfs(SurfaceMeta *seed, const std::vector &surfs_v, const nlohmann::json ¶ms, float voxelsize = 1.0); diff --git a/core/src/GrowPatch.cpp b/core/src/GrowPatch.cpp index 8d978d6d7..6cc2f8a1d 100644 --- a/core/src/GrowPatch.cpp +++ b/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); } From 918e6aa6270b9d5f150056efab5582978f07ce08 Mon Sep 17 00:00:00 2001 From: Paul Henderson Date: Fri, 5 Sep 2025 14:32:08 +0100 Subject: [PATCH 2/2] Adjust normalisation for field weight volumes to be more intuitive --- core/include/vc/core/types/ChunkedTensor.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/include/vc/core/types/ChunkedTensor.hpp b/core/include/vc/core/types/ChunkedTensor.hpp index c40a2d158..51b8b4637 100644 --- a/core/include/vc/core/types/ChunkedTensor.hpp +++ b/core/include/vc/core/types/ChunkedTensor.hpp @@ -639,7 +639,7 @@ struct Chunked3dFloatFromUint8 p *= _scale; cv::Vec3i i{lround(p[0]), lround(p[1]), lround(p[2])}; uint8_t x = _x(i) ; - return (float{x} - 128.f) / 127.f; + return float{x} / 255.f; } float operator()(double z, double y, double x)