Skip to content

Commit

Permalink
fix: closes #32
Browse files Browse the repository at this point in the history
  • Loading branch information
@kdmarrett
kdmarrett committed Oct 14, 2023
1 parent 56b3f42 commit 958921c
Showing 1 changed file with 72 additions and 52 deletions.
124 changes: 72 additions & 52 deletions src/mesh_reconstruction.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
#include <GEL/Geometry/graph_io.h>
#include <GEL/Geometry/graph_util.h>
#include <GEL/Geometry/graph_skeletonize.h>
//#include <GEL/Geometry/graph_util.h>
#include <GEL/Geometry/KDTree.h>
#include <GEL/Util/AttribVec.h>
#include <openvdb/tools/FastSweeping.h> // fogToSdf
Expand All @@ -14,6 +13,29 @@
#include <openvdb/tools/ValueTransformer.h>
#include <openvdb/tools/VolumeToMesh.h>

class Node {
public:
CGLA::Vec3d pos;
float radius;
};

using NodeID = Geometry::AMGraph3D::NodeID;

float get_radius(Util::AttribVec<NodeID, CGLA::Vec3f> &node_color, NodeID 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};
}

Node get_node(Geometry::AMGraph3D g, NodeID i) {
Node n{g.pos[i], get_radius(g.node_color, i)};
return n;
}

auto euc_dist = [](auto a, auto b) -> float {
std::array<float, 3> diff = {
static_cast<float>(a[0]) - static_cast<float>(b[0]),
Expand All @@ -23,8 +45,8 @@ auto euc_dist = [](auto a, auto b) -> float {
};

// build kdtree, highly efficient for nearest point computations
Geometry::KDTree<CGLA::Vec3d, unsigned long> build_kdtree(Geometry::AMGraph3D &graph) {
Geometry::KDTree<CGLA::Vec3d, unsigned long> tree;
Geometry::KDTree<CGLA::Vec3d, NodeID> build_kdtree(Geometry::AMGraph3D &graph) {
Geometry::KDTree<CGLA::Vec3d, NodeID> tree;
for (auto i : graph.node_ids()) {
CGLA::Vec3d p0 = graph.pos[i];
tree.insert(p0, i);
Expand All @@ -34,34 +56,34 @@ 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) {
std::vector<NodeID> within_sphere(Node &node,
Geometry::KDTree<CGLA::Vec3d, NodeID> &tree,
float soma_dilation=1) {

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 * seed.radius, keys, vals);
std::vector<CGLA::Vec3d> _;
std::vector<NodeID> vals;
tree.in_sphere(node.pos, soma_dilation * node.radius, _, vals);
return vals;
}

// 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
// so you may need to artificially set the radius to something small then
// merge an ideal sphere after meshing
// in use cases where the graph is not transformed into a mesh this is
// not necessary
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);

std::set<unsigned long> deleted_nodes;
std::set<NodeID> deleted_nodes;
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 = CGLA::Vec3d(seed.coord[0], seed.coord[1], seed.coord[2]);
Node n{original_pos, static_cast<float>(seed.radius)};

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

auto new_seed_id = graph.merge_nodes(within_sphere_ids);
Expand Down Expand Up @@ -94,8 +116,8 @@ std::vector<GridCoord> force_soma_nodes(Geometry::AMGraph3D &graph,
return soma_coords;
}

void check_soma_ids(unsigned long nodes, std::vector<unsigned long> soma_ids) {
rng::for_each(soma_ids, [&](unsigned long soma_id) {
void check_soma_ids(NodeID nodes, std::vector<NodeID> soma_ids) {
rng::for_each(soma_ids, [&](NodeID soma_id) {
if (soma_id >= nodes) {
throw std::runtime_error("Impossible soma id found");
}
Expand All @@ -109,10 +131,12 @@ std::vector<GridCoord> find_soma_nodes(Geometry::AMGraph3D &graph,

std::vector<GridCoord> soma_coords;
rng::for_each(seeds, [&](Seed seed) {
std::optional<size_t> max_index;
std::optional<NodeID> max_index;

if (highest_valence) {
auto within_sphere_ids = within_sphere(seed, tree, soma_dilation);
auto original_pos = CGLA::Vec3d(seed.coord[0], seed.coord[1], seed.coord[2]);
Node n{original_pos, static_cast<float>(seed.radius)};
auto within_sphere_ids = within_sphere(n, tree, soma_dilation);

// pick skeletal node within radii with highest number of edges (valence)
int max_valence = 0;
Expand All @@ -127,7 +151,7 @@ std::vector<GridCoord> find_soma_nodes(Geometry::AMGraph3D &graph,
auto coord = seed.coord;
CGLA::Vec3d p0(coord[0], coord[1], coord[2]);
CGLA::Vec3d key;
unsigned long val;
NodeID val;
double dist = 1000;
bool found = tree.closest_point(p0, dist, key, val);
if (found)
Expand Down Expand Up @@ -187,16 +211,6 @@ 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 Down Expand Up @@ -420,15 +434,26 @@ 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) {
void test_all_valid_radii(Geometry::AMGraph3D &g) {
for (NodeID 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));
}
}
}

void test_no_node_within_another(Geometry::AMGraph3D &g) {
auto tree = build_kdtree(g);
for (NodeID i=0; i < g.no_nodes(); ++i) {
auto n = get_node(g, i);
std::vector<NodeID> within_sphere_ids = within_sphere(n, tree);
if (within_sphere_ids.size() > 1) {
throw std::runtime_error("Found node within another");
}
}
}

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 @@ -467,18 +492,13 @@ 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";
//component_log << "soma nodes\n";
// sweep through various soma ids
std::vector<GridCoord> soma_coords;
if (args->seed_action == "force")
Expand All @@ -488,14 +508,14 @@ 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);
//test_no_node_within_another(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";
//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);
Expand Down Expand Up @@ -548,11 +568,11 @@ void write_swcs(Geometry::AMGraph3D component_graph, std::vector<GridCoord> soma

// each vertex in the graph has a single parent (id) which is
// determined via BFS traversal
std::vector<int> parent_table(component_graph.no_nodes(), -1);
std::vector<NodeID> parent_table(component_graph.no_nodes(), -1);

// scan the graph to find the final set of soma_ids
std::vector<unsigned long> soma_ids;
for (int i=0; i < component_graph.no_nodes(); ++i) {
std::vector<NodeID> soma_ids;
for (NodeID i=0; i < component_graph.no_nodes(); ++i) {
auto pos = component_graph.pos[i];
auto current_coord = GridCoord(pos[0], pos[1], pos[2]);
if (std::find(soma_coords.begin(),
Expand All @@ -563,15 +583,15 @@ void write_swcs(Geometry::AMGraph3D component_graph, std::vector<GridCoord> soma
// do BFS from each known soma in the component
for (auto soma_id : soma_ids) {
// init q with soma
std::queue<size_t> q;
std::queue<NodeID> q;
q.push(soma_id);

// start swc and add header metadata
auto pos = component_graph.pos[soma_id].get();
auto file_name_base = "tree-with-soma-xyz-" +
std::to_string((int)pos[0]) + "-" +
std::to_string((int)pos[1]) + "-" +
std::to_string((int)pos[2]);
std::to_string((NodeID)pos[0]) + "-" +
std::to_string((NodeID)pos[1]) + "-" +
std::to_string((NodeID)pos[2]);
auto coord_to_swc_id = get_id_map();

// traverse rest of tree
Expand All @@ -597,7 +617,7 @@ void write_swcs(Geometry::AMGraph3D component_graph, std::vector<GridCoord> soma
}

while (q.size()) {
size_t id = q.front();
NodeID id = q.front();
q.pop();

auto pos = component_graph.pos[id].get();
Expand Down Expand Up @@ -648,13 +668,13 @@ swc_to_graph(filesystem::path marker_file, std::array<double, 3> voxel_size,
auto min_voxel_size = min_max(voxel_size).first;
std::vector<Seed> seeds;
Geometry::AMGraph3D g;
int count = 0;
NodeID count = 0;
while (ifs.good()) {
if (ifs.peek() == '#' || ifs.eof()) {
ifs.ignore(1000, '\n');
continue;
}
int id, type, parent_id;
NodeID id, type, parent_id;
double x_um,y_um,z_um,radius_um;
ifs >> id;
ifs.ignore(10, ' ');
Expand Down Expand Up @@ -712,7 +732,7 @@ openvdb::FloatGrid::Ptr skeleton_to_surface(Geometry::AMGraph3D skeleton) {
// feq requires an explicit radii vector
std::vector<double> radii;
radii.reserve(skeleton.no_nodes());
for (int i = 0; i < skeleton.no_nodes(); ++i) {
for (NodeID i = 0; i < skeleton.no_nodes(); ++i) {
//std::cout << get_radius(skeleton.node_color, i) << '\n';
radii.push_back(get_radius(skeleton.node_color, i));
}
Expand All @@ -722,7 +742,7 @@ 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);
HMesh::obj_save("mesh.obj", manifold);

// translate manifold into something vdb will understand
// points ands quads
Expand Down

0 comments on commit 958921c

Please sign in to comment.