Skip to content

Commit

Permalink
fix: close #30
Browse files Browse the repository at this point in the history
  • Loading branch information
@kdmarrett
kdmarrett committed Oct 13, 2023
1 parent c6d3876 commit 56b3f42
Show file tree
Hide file tree
Showing 8 changed files with 130 additions and 102 deletions.
176 changes: 98 additions & 78 deletions src/mesh_reconstruction.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,48 +36,40 @@ Geometry::KDTree<CGLA::Vec3d, unsigned long> build_kdtree(Geometry::AMGraph3D &g
// find existing skeletal node within the radius of the soma
std::vector<unsigned long> within_sphere(Seed &seed,
Geometry::KDTree<CGLA::Vec3d, unsigned long> &tree,
float soma_dilation, bool use_voxel_units = false) {
float soma_dilation) {

CGLA::Vec3d p0 = use_voxel_units ? CGLA::Vec3d(seed.coord[0], seed.coord[1],
seed.coord[2]) : CGLA::Vec3d(seed.coord_um[0],
seed.coord_um[1], seed.coord_um[2]);
double radius = use_voxel_units ? seed.radius : seed.radius_um;
CGLA::Vec3d p0{static_cast<double>(seed.coord[0]),
static_cast<double>(seed.coord[1]), static_cast<double>(seed.coord[2])};

std::vector<CGLA::Vec3d> keys;
std::vector<unsigned long> vals;
tree.in_sphere(p0, soma_dilation * radius, keys, vals);
tree.in_sphere(p0, soma_dilation * seed.radius, keys, vals);
return vals;
}

void merge_local_radius(Geometry::AMGraph3D &graph, std::vector<std::pair<Seed,
unsigned long>> seed_pairs, float soma_dilation=1., bool use_voxel_units = false) {
// nodes with high valence and large radius have a bug with FEQ remesh
// so you need to artificially set the radius to something small then
// merge an ideal sphere after meshing
void merge_local_radius(Geometry::AMGraph3D &graph, std::vector<Seed> &seeds,
float soma_dilation=1., bool keep_radius_small=false) {
// the graph is complete so build a data structure that
// is fast at finding nodes within a 3D radial distance
auto tree = build_kdtree(graph);

auto soma_positions =
seed_pairs | rv::transform([&](auto seedp) {
auto [seed, _] = seedp;
return use_voxel_units ? CGLA::Vec3d(seed.coord[0], seed.coord[1],
seed.coord[2]) : CGLA::Vec3d(seed.coord_um[0],
seed.coord_um[1], seed.coord_um[2]); })
| rng::to_vector;

std::set<unsigned long> deleted_nodes;
rng::for_each(seed_pairs, [&](auto seedp) {
auto [seed, _] = seedp;
rng::for_each(seeds, [&](Seed seed) {
//std::cout << "Seed " << seed.coord_um[0] << ' ' << seed.coord_um[1] << ' ' << seed.coord_um[2] << " radius " << seed.radius_um << '\n';
auto original_pos = use_voxel_units ? CGLA::Vec3d(seed.coord[0], seed.coord[1],
seed.coord[2]) : CGLA::Vec3d(seed.coord_um[0],
seed.coord_um[1], seed.coord_um[2]);
auto original_pos = CGLA::Vec3d(seed.coord[0], seed.coord[1], seed.coord[2]);

std::vector<unsigned long> within_sphere_ids = within_sphere(seed, tree, soma_dilation, use_voxel_units);
std::vector<unsigned long> within_sphere_ids = within_sphere(seed, tree, soma_dilation);
//std::cout << "removing " << (within_sphere_ids.size() - 1) << '\n';

auto new_seed_id = graph.merge_nodes(within_sphere_ids);

// set the position and radius explicitly rather than averaging the merged nodes
graph.pos[new_seed_id] = original_pos;
graph.node_color[new_seed_id] = CGLA::Vec3f(0, use_voxel_units ? seed.radius : seed.radius_um, 0);

graph.node_color[new_seed_id] = CGLA::Vec3f(0, keep_radius_small ? 1 : seed.radius, 0);
});

graph = Geometry::clean_graph(graph);
Expand All @@ -92,26 +84,12 @@ std::vector<GridCoord> force_soma_nodes(Geometry::AMGraph3D &graph,
std::vector<Seed> &seeds,
float soma_dilation) {

// add the known seeds to the skeletonized graph
// aggregate their seed ids so they are not mistakenly merged
// below
auto soma_coords =
seeds | rv::transform([](Seed seed) { return seed.coord; })
| rng::to_vector;

auto new_soma_ids =
seeds | rv::transform([&](Seed seed) {
auto p = CGLA::Vec3d(seed.coord[0], seed.coord[1], seed.coord[2]);
auto soma_id = graph.add_node(p);
// set radius by previously computed radius
graph.node_color[soma_id] =
CGLA::Vec3f(0, soma_dilation * seed.radius, 0);
return soma_id;
}) | rng::to_vector;

auto seed_pairs = rv::zip(seeds, new_soma_ids) | rng::to_vector;

merge_local_radius(graph, seed_pairs, soma_dilation, /*voxelunits*/true);
// add the known seeds to the skeletonized graph
merge_local_radius(graph, seeds, soma_dilation);

return soma_coords;
}
Expand All @@ -134,7 +112,7 @@ std::vector<GridCoord> find_soma_nodes(Geometry::AMGraph3D &graph,
std::optional<size_t> max_index;

if (highest_valence) {
auto within_sphere_ids = within_sphere(seed, tree, soma_dilation, true);
auto within_sphere_ids = within_sphere(seed, tree, soma_dilation);

// pick skeletal node within radii with highest number of edges (valence)
int max_valence = 0;
Expand Down Expand Up @@ -209,6 +187,16 @@ Geometry::AMGraph3D vdb_to_graph(openvdb::FloatGrid::Ptr component,
return g;
}

float get_radius(Util::AttribVec<Geometry::AMGraph::NodeID, CGLA::Vec3f> &node_color, size_t i) {
auto color = node_color[i].get();
// radius is in the green channel
return color[1]; // RGB
}

CGLA::Vec3f convert_radius(float radius) {
return {0, radius, 0};
}

// naive multifurcation fix, force non-soma vertices to have at max 3 neighbors
Geometry::AMGraph3D fix_multifurcations(Geometry::AMGraph3D &graph,
std::vector<GridCoord> soma_coords) {
Expand All @@ -227,8 +215,11 @@ Geometry::AMGraph3D fix_multifurcations(Geometry::AMGraph3D &graph,
// build a new averaged node
CGLA::Vec3d pos1 = graph.pos[multifurc_id];
CGLA::Vec3d pos2 = graph.pos[to_extend];
auto rad = (get_radius(graph.node_color, multifurc_id) +
get_radius(graph.node_color, to_extend)) / 2;
auto pos3 = CGLA::Vec3d((pos1[0] + pos2[0]) / 2, (pos1[1] + pos2[1]) / 2, (pos1[2] + pos2[2]) / 2);
auto new_path_node = graph.add_node(pos3);
graph.node_color[new_path_node] = convert_radius(rad);

graph.connect_nodes(new_path_node, to_extend);
graph.connect_nodes(new_path_node, multifurc_id);
Expand Down Expand Up @@ -369,12 +360,6 @@ openvdb::FloatGrid::Ptr mask_to_sdf(openvdb::MaskGrid::Ptr mask) {
return vto::fogToSdf(*float_grid, 0);
}

float get_radius(Util::AttribVec<Geometry::AMGraph::NodeID, CGLA::Vec3f> &node_color, size_t i) {
auto color = node_color[i].get();
// radius is in the green channel
return color[1]; // RGB
}

// laplacian smoothing, controllable by number of repeated iterations
// and the smoothing strength (alpha) at each iteration
// a high number of iterations allows the smoothing effect to diffuse
Expand All @@ -385,10 +370,6 @@ float get_radius(Util::AttribVec<Geometry::AMGraph::NodeID, CGLA::Vec3f> &node_c
// after only 1 iteration at 1 alpha
void smooth_graph_pos_rad(Geometry::AMGraph3D &g, const int iter, const float alpha) {
auto lsmooth = [](Geometry::AMGraph3D &g, float _alpha) {
auto convert_radius = [](float radius) {
return CGLA::Vec3f(0, radius, 0);
};

Util::AttribVec<Geometry::AMGraph::NodeID, CGLA::Vec3d> new_pos(
g.no_nodes(), CGLA::Vec3d(0));
Util::AttribVec<Geometry::AMGraph::NodeID, CGLA::Vec3f> new_radius(
Expand Down Expand Up @@ -439,6 +420,15 @@ void smooth_graph_pos_rad(Geometry::AMGraph3D &g, const int iter, const float al
} }
*/

void test_all_valid_radii(Geometry::AMGraph3D g) {
for (long int i=0; i < g.no_nodes(); ++i) {
auto rad = get_radius(g.node_color, i);
if (rad < .001) {
throw std::runtime_error("Found node with radii " + std::to_string(rad));
}
}
}

std::optional<std::pair<Geometry::AMGraph3D, std::vector<GridCoord>>>
vdb_to_skeleton(openvdb::FloatGrid::Ptr component, std::vector<Seed> component_seeds,
int index, RecutCommandLineArgs *args,
Expand Down Expand Up @@ -477,12 +467,18 @@ vdb_to_skeleton(openvdb::FloatGrid::Ptr component, std::vector<Seed> component_s
auto [component_graph, _] = skeleton_from_node_set_vec(g, separators);
component_log << "msls, " << timer.elapsed() << '\n';

test_all_valid_radii(component_graph);
// prune all leaf vertices (valency 1) whose only neighbor has valency > 2
// as these tend to be spurious branches
component_log << "prune\n";
Geometry::prune(component_graph);
test_all_valid_radii(component_graph);

component_log << "prune\n";
smooth_graph_pos_rad(component_graph, args->smooth_steps, /*alpha*/ 1);
test_all_valid_radii(component_graph);

component_log << "soma nodes\n";
// sweep through various soma ids
std::vector<GridCoord> soma_coords;
if (args->seed_action == "force")
Expand All @@ -492,13 +488,17 @@ vdb_to_skeleton(openvdb::FloatGrid::Ptr component, std::vector<Seed> component_s
else if (args->seed_action == "find-valent")
soma_coords = find_soma_nodes(component_graph, component_seeds, args->soma_dilation.value(), true);

test_all_valid_radii(component_graph);

if (soma_coords.size() == 0) {
std::cerr << "Warning no soma_coords found for component " << index << '\n';
return std::nullopt;
}

component_log << "multifurcations\n";
// multifurcations are only important for rules of SWC standard
component_graph = fix_multifurcations(component_graph, soma_coords);
test_all_valid_radii(component_graph);

if (save_graphs)
graph_save(component_dir_fn / ("skeleton.graph"), component_graph);
Expand All @@ -515,7 +515,7 @@ void write_ano_file(fs::path component_dir_fn, std::string file_name_base) {
}

void write_apo_file(fs::path component_dir_fn, std::string file_name_base, std::array<double, 3> pos,
float radius, std::array<float, 3> voxel_size) {
float radius, std::array<double, 3> voxel_size) {
std::ofstream apo_file;
apo_file.open(component_dir_fn / (file_name_base + ".ano.apo"));
apo_file << std::fixed << std::setprecision(SWC_PRECISION);
Expand All @@ -541,7 +541,7 @@ void write_apo_file(fs::path component_dir_fn, std::string file_name_base, std::
}

void write_swcs(Geometry::AMGraph3D component_graph, std::vector<GridCoord> soma_coords,
std::array<float, 3> voxel_size,
std::array<double, 3> voxel_size,
std::filesystem::path component_dir_fn = ".",
CoordBBox bbox = {}, bool bbox_adjust = false,
bool is_eswc = false, bool disable_swc_scaling = false) {
Expand Down Expand Up @@ -637,15 +637,16 @@ void write_swcs(Geometry::AMGraph3D component_graph, std::vector<GridCoord> soma
//}
}

std::pair<Geometry::AMGraph3D, std::vector<std::pair<Seed, unsigned long>>>
swc_to_graph(filesystem::path marker_file, bool save_file = false) {
std::pair<Geometry::AMGraph3D, std::vector<Seed>>
swc_to_graph(filesystem::path marker_file, std::array<double, 3> voxel_size,
bool save_file = false) {
ifstream ifs(marker_file);
if (ifs.fail()) {
throw std::runtime_error("Unable to open marker file " + marker_file.string());
}

auto min_voxel_size = min_max(voxel_size).first;
std::vector<Seed> seeds;
std::vector<unsigned long> seed_ids;
Geometry::AMGraph3D g;
int count = 0;
while (ifs.good()) {
Expand All @@ -654,45 +655,56 @@ swc_to_graph(filesystem::path marker_file, bool save_file = false) {
continue;
}
int id, type, parent_id;
double x,y,z,radius_um;
double x_um,y_um,z_um,radius_um;
ifs >> id;
ifs.ignore(10, ' ');
ifs >> type;
ifs.ignore(10, ' ');
ifs >> x;
ifs >> x_um;
ifs.ignore(10, ' ');
ifs >> y;
ifs >> y_um;
ifs.ignore(10, ' ');
ifs >> z;
ifs >> z_um;
ifs.ignore(10, ' ');
ifs >> radius_um;
ifs.ignore(10, ' ');
ifs >> parent_id;
ifs.ignore(1000, '\n');
parent_id -= 1; // need to adjust to 0-indexed

// translate from um units into integer voxel units
double x,y,z,radius;
radius = radius_um / min_voxel_size;
x = x_um / voxel_size[0];
y = y_um / voxel_size[1];
z = z_um / voxel_size[2];
auto p = CGLA::Vec3d(x, y, z);
auto coord = GridCoord(x, y, z);

// add it to the graph
auto node_id = g.add_node(p);
g.node_color[node_id] = CGLA::Vec3f(0, radius_um, 0);
//std::cout << x << ' ' << y << ' ' << z << ' ' << radius_um << ' ' << parent_id << '\n';
g.node_color[node_id] = CGLA::Vec3f(0, radius, 0);
//std::cout << x << ' ' << y << ' ' << z << ' ' << radius << ' ' << parent_id << '\n';

// somas are nodes that have a parent of -1 (adjusted to -2) or have an index
// of themselves
if (parent_id == -2 || parent_id == node_id) {
std::array<double, 3> coord_um{x, y, z};
seeds.emplace_back(coord_um, radius_um);
seed_ids.push_back(id);
auto volume = static_cast<uint64_t>((4 / 3) * PI * std::pow(radius, 3));
std::array<double, 3> coord_um{p[0], p[1], p[2]};
seeds.emplace_back(coord, coord_um, radius, radius_um, volume);
} else {
g.connect_nodes(node_id, parent_id);
}
}

auto seed_pairs = rv::zip(seeds, seed_ids) | rng::to_vector;

if (seeds.size() > 1) {
throw std::runtime_error("Warning: SWC files are trees which by definition must have only 1 root (soma), provided file has " + seeds.size());
}

if (save_file)
graph_save(marker_file.stem().string() + ".graph", g);

return std::make_pair(g, seed_pairs);
return std::make_pair(g, seeds);
}

openvdb::FloatGrid::Ptr skeleton_to_surface(Geometry::AMGraph3D skeleton) {
Expand All @@ -710,6 +722,8 @@ openvdb::FloatGrid::Ptr skeleton_to_surface(Geometry::AMGraph3D skeleton) {
std::cout << "Start graph to manifold" << '\n';
auto manifold = graph_to_FEQ(skeleton, radii);

//HMesh::obj_save("mesh.obj", manifold);

// translate manifold into something vdb will understand
// points ands quads
std::cout << "Start manifold to polygons" << '\n';
Expand All @@ -722,32 +736,38 @@ openvdb::FloatGrid::Ptr skeleton_to_surface(Geometry::AMGraph3D skeleton) {
}

// returns a polygonal mesh
openvdb::FloatGrid::Ptr swc_to_segmented(filesystem::path marker_file, bool save_vdbs = false,
openvdb::FloatGrid::Ptr swc_to_segmented(filesystem::path marker_file,
std::array<double, 3> voxel_size, bool save_vdbs = false,
std::string name = "") {
auto [skeleton, seed_pairs] = swc_to_graph(marker_file);
auto [skeleton, seeds] = swc_to_graph(marker_file, voxel_size);

merge_local_radius(skeleton, seed_pairs);
merge_local_radius(skeleton, seeds, 1, true);

if (save_vdbs)
graph_save(marker_file.stem().string() + ".graph", skeleton);

auto seeds = seed_pairs | rv::keys | rng::to_vector;
auto soma_coords = rv::transform(seeds, [](Seed seed) {
return GridCoord(seed.coord_um[0], seed.coord_um[1], seed.coord_um[2]);
}) | rng::to_vector;

write_swcs(skeleton, soma_coords, {1,1,1}, "test");
if (false) {
auto soma_coords = rv::transform(seeds, [](Seed seed) {
return seed.coord;
}) | rng::to_vector;
write_swcs(skeleton, soma_coords, voxel_size, "test");
}

auto level_set = skeleton_to_surface(skeleton);

// merge soma on top
{
auto [_, merged_somas] = create_seed_sphere_grid(seeds);
csgUnion(*level_set, *merged_somas);
// only 1 soma allowed per swc
auto seed = seeds[0];
auto soma_sdf = vto::createLevelSetSphere<openvdb::FloatGrid>(
seed.radius, seed.coord.asVec3s(), 1.,
RECUT_LEVEL_SET_HALF_WIDTH);
vto::csgUnion(*level_set, *soma_sdf); // empties merged_somas grid into ls
}

if (save_vdbs)
write_vdb_file({level_set}, "surface-" + (name.empty() ? marker_file.stem().string() : name) + ".vdb");

return level_set;
}

Expand Down
Loading

0 comments on commit 56b3f42

Please sign in to comment.