Pre labelled tree extraction

raycloudtools/rayextract/rayextract.cpp

@@ -62,6 +62,7 @@ void usage(int exit_code = 1)
     std::cout << "                            --branch_segmentation- (-b) _segmented.ply is per branch segment" << std::endl;
     std::cout << "                            --grid_width 10      - (-w) crops results assuming cloud has been gridded with given width" << std::endl;
     std::cout << "                            --use_rays           - (-u) use rays to reduce trunk radius overestimation in noisy cloud data" << std::endl;
+    std::cout << "                            --segmented          - (-l) the cloud is already segmented into trees, one per colour (as in _segmented.ply). Only the tree models are estimated, black points are unlabelled" << std::endl;
     std::cout << "                            (for internal constants -c -g -s see source file rayextract)" << std::endl;
     std::cout << "   note: branch radii rely on priors (Leonardo's rule at splits, self-similar taper) that can bias branch" << std::endl;
     std::cout << "         measurements. Only use these results if the mechanism and its implications are understood." << std::endl;
@@ -100,6 +101,7 @@ int rayExtract(int argc, char *argv[])
   ray::DoubleArgument gradient(0.001, 1000.0, 1.0), global_taper(0.0, 1.0), global_taper_factor(0.0, 1.0);
   ray::OptionalKeyValueArgument gradient_option("gradient", 'g', &gradient);
   ray::OptionalFlagArgument exclude_rays("exclude_rays", 'e'), segment_branches("branch_segmentation", 'b'), stalks("stalks", 's'), use_rays("use_rays", 'u');
+  ray::OptionalFlagArgument segmented("segmented", 'l');
   ray::DoubleArgument width(0.01, 10.0, 0.25), drop(0.001, 1.0), max_gradient(0.01, 5.0), min_gradient(0.01, 5.0);
 
   ray::DoubleArgument max_diameter(0.01, 100.0), distance_limit(0.01, 10.0), height_min(0.01, 1000.0),
@@ -140,7 +142,8 @@ int rayExtract(int argc, char *argv[])
     argc, argv, { &trees, &cloud_file, &mesh_file },
     { &max_diameter_option, &distance_limit_option, &height_min_option, &crop_length_option, &girth_height_ratio_option,
       &cylinder_length_to_width_option, &gap_ratio_option, &span_ratio_option, &gravity_factor_option,
-      &segment_branches, &grid_width_option, &global_taper_option, &global_taper_factor_option, &use_rays, &verbose, &view_flag });
+      &segment_branches, &grid_width_option, &global_taper_option, &global_taper_factor_option, &use_rays, &segmented,
+      &verbose, &view_flag });
   bool extract_leaves = ray::parseCommandLine(argc, argv, { &leaves, &cloud_file, &trees_file }, { &leaf_option, &leaf_area_option, &leaf_droop_option, &stalks, &view_flag });
 
   if (!extract_trunks && !extract_forest && !extract_terrain && !extract_trees && !extract_leaves)
@@ -233,7 +236,34 @@ int rayExtract(int argc, char *argv[])
     params.use_rays = use_rays.isSet(); 
     params.segment_branches = segment_branches.isSet();
 
-    ray::Trees trees(cloud, offset, mesh, params, verbose.isSet());
+    // for a pre-segmented cloud we take the tree label of each point from its colour, so the tree
+    // structures are estimated per label rather than segmenting the trees from the point data
+    std::vector<int> point_labels;
+    if (segmented.isSet())
+    {
+      point_labels.reserve(cloud.rayCount());
+      size_t num_labelled = 0;
+      for (size_t i = 0; i < cloud.rayCount(); i++)
+      {
+        // this is the inverse of the colouring used when saving the _segmented.ply file, and as it is a
+        // one-to-one mapping it also works for clouds coloured one arbitrary colour per tree.
+        // black (unlabelled) points convert to -1
+        point_labels.push_back(ray::convertColourToInt(cloud.colours[i]));
+        if (point_labels.back() >= 0)
+        {
+          num_labelled++;
+        }
+      }
+      if (num_labelled == 0)
+      {
+        std::cerr << "Error: --segmented expects one colour per tree in " << cloud_file.name()
+                  << ", but every point is black (unlabelled)" << std::endl;
+        usage(true);
+      }
+    }
+
+    ray::Trees trees(cloud, offset, mesh, params, verbose.isSet(),
+                     segmented.isSet() ? &point_labels : nullptr);
 
     // output the picewise cylindrical description of the trees
     trees.save(cloud_file.nameStub() + "_trees.txt", offset, verbose.isSet());

raylib/extraction/raysegment.cpp

@@ -6,7 +6,10 @@
 #include "raysegment.h"
 #include <nabo/nabo.h>
 #include "rayterrain.h"
+#include <iostream>
+#include <map>
 #include <queue>
+#include <unordered_map>
 
 namespace ray
 {
@@ -40,10 +43,12 @@ class QueueNodeComparator
 /// @c distance_limit maximum distance between points that can be connected
 /// @c gravity_factor controls how far laterally the shortest paths can travel
 /// @c closest_node a priority queue
+/// @c path_labels optional per-point tree labels (-1 for unlabelled). When supplied, paths are not allowed to
+/// connect points of two different labels, and unlabelled points adopt the label of the path that reaches them
 void connectPointsShortestPath(
   std::vector<Vertex> &points,
   std::priority_queue<QueueNode, std::vector<QueueNode>, QueueNodeComparator> &closest_node, double distance_limit,
-  double gravity_factor)
+  double gravity_factor, std::vector<int> *path_labels)
 {
   // 1. get nearest neighbours
   const int search_size = std::min(20, static_cast<int>(points.size()) - 1);
@@ -71,6 +76,16 @@ void connectPointsShortestPath(
       for (int i = 0; i < search_size && indices(i, node.id) != Nabo::NNSearchD::InvalidIndex; i++)
       {
         const int child = indices(i, node.id);
+        if (path_labels)
+        {
+          // paths are not allowed to cross from one labelled tree into another
+          const int node_label = (*path_labels)[node.id];
+          const int child_label = (*path_labels)[child];
+          if (node_label != -1 && child_label != -1 && node_label != child_label)
+          {
+            continue;
+          }
+        }
         const double dist2 = dists2(i, node.id);  // square distance to neighbour
         const double dist = std::sqrt(dist2);
         double new_score = 0;
@@ -117,6 +132,11 @@ void connectPointsShortestPath(
           points[child].distance_to_ground = node.distance_to_ground + dist;
           points[child].parent = node.id;
           points[child].root = node.root;
+          if (path_labels && (*path_labels)[child] == -1)
+          {
+            // an unlabelled point belongs to the tree that its path comes from
+            (*path_labels)[child] = (*path_labels)[node.id];
+          }
           closest_node.push(
             QueueNode(points[child].distance_to_ground, points[child].score, node.radius, node.root, child));
         }
@@ -130,17 +150,56 @@ void connectPointsShortestPath(
 /// the returned vector of index sets provides the root points for each separated tree
 std::vector<std::vector<int>> getRootsAndSegment(std::vector<Vertex> &points, const Cloud &cloud, const Mesh &mesh,
                                                  double max_diameter, double distance_limit, double height_min,
-                                                 double gravity_factor)
+                                                 double gravity_factor, const std::vector<int> *point_labels)
 {
+  if (point_labels && point_labels->size() != cloud.ends.size())
+  {
+    std::cerr << "Error: the number of point labels does not match the number of rays in the cloud" << std::endl;
+    return std::vector<std::vector<int>>();
+  }
   // first fill in the basic attributes of the points structure
   points.reserve(cloud.ends.size());
+  std::vector<int> labels;  // the supplied label per point, or -1 where unlabelled
+  int num_labels = 0;
+  if (point_labels)
+  {
+    labels.reserve(cloud.ends.size());
+  }
+  // maps the supplied label values onto a contiguous set of tree indices, in ascending label order so that
+  // the output trees keep the order (and so the IDs) of the input labels
+  std::map<int, int> label_ids;
+  if (point_labels)
+  {
+    for (unsigned int i = 0; i < cloud.ends.size(); i++)
+    {
+      if (cloud.rayBounded(i) && (*point_labels)[i] >= 0)
+      {
+        label_ids[(*point_labels)[i]] = 0;
+      }
+    }
+    for (auto &label_id : label_ids)
+    {
+      label_id.second = num_labels++;
+    }
+  }
   for (unsigned int i = 0; i < cloud.ends.size(); i++)
   {
     if (cloud.rayBounded(i))
     {
       points.push_back(Vertex(cloud.ends[i], cloud.starts[i]));
+      if (point_labels)
+      {
+        // negative labels are unlabelled, so these points attach to whichever tree's path reaches them
+        const int label = (*point_labels)[i];
+        labels.push_back(label < 0 ? -1 : label_ids[label]);
+      }
     }
   }
+  const int num_cloud_points = static_cast<int>(points.size());
+  if (point_labels)
+  {
+    std::cout << "using " << num_labels << " pre-labelled trees from the input cloud" << std::endl;
+  }
 
   const double pixel_width = max_diameter;
   Eigen::Vector3d box_min, box_max;
@@ -157,6 +216,13 @@ std::vector<std::vector<int>> getRootsAndSegment(std::vector<Vertex> &points, co
       points.push_back(Vertex(vert - Eigen::Vector3d(0,0,0.01), vert + Eigen::Vector3d(0,0,0.01))); // make small vertical ray
     }
   }
+  // the ground mesh vertices are unlabelled, they are the seeds that the labelled trees grow from
+  std::vector<int> path_labels;
+  if (point_labels)
+  {
+    path_labels = labels;
+    path_labels.resize(points.size(), -1);
+  }
   // convert the ground mesh to an easy look-up height field
   Eigen::ArrayXXd lowfield;
   mesh.toHeightField(lowfield, box_min, box_max, pixel_width);
@@ -190,7 +256,120 @@ std::vector<std::vector<int>> getRootsAndSegment(std::vector<Vertex> &points, co
   }
 
   // perform Djikstra's shortest path to ground algorithm to fill in the parent indices in 'points'
-  connectPointsShortestPath(points, closest_node, distance_limit, gravity_factor);
+  connectPointsShortestPath(points, closest_node, distance_limit, gravity_factor,
+                            point_labels ? &path_labels : nullptr);
+
+  // when labels have been supplied we don't need to segment the trees, we just group the root points
+  // according to the label of the points that grow from them, generating one tree per label
+  if (point_labels)
+  {
+    // each labelled point votes for its root point, so that a root point shared by several
+    // labels is given to the label that most of its points belong to
+    std::vector<std::unordered_map<int, int>> votes(points.size() - roots_start);
+    std::vector<double> max_heights(num_labels, 0.0);
+    for (int i = 0; i < num_cloud_points; i++)
+    {
+      const int label = labels[i];
+      if (label == -1)
+      {
+        continue;
+      }
+      const Eigen::Vector3i index = ((points[i].pos - box_min) / pixel_width).cast<int>();
+      max_heights[label] = std::max(max_heights[label], points[i].pos[2] - lowfield(index[0], index[1]));
+      if (points[i].root >= roots_start)
+      {
+        votes[points[i].root - roots_start][label]++;
+      }
+    }
+    std::vector<std::vector<int>> roots_set(num_labels);
+    std::vector<int> best_root(num_labels, -1);       // the most voted for root point per label
+    std::vector<int> best_root_votes(num_labels, 0);
+    for (int i = roots_start; i < static_cast<int>(points.size()); i++)
+    {
+      int best_label = -1;
+      int most_votes = 0;
+      for (const auto &vote : votes[i - roots_start])
+      {
+        if (vote.second > most_votes)
+        {
+          most_votes = vote.second;
+          best_label = vote.first;
+        }
+      }
+      if (best_label != -1)
+      {
+        roots_set[best_label].push_back(i);
+        if (most_votes > best_root_votes[best_label])
+        {
+          best_root_votes[best_label] = most_votes;
+          best_root[best_label] = i;
+        }
+      }
+    }
+    // only keep the root points close to the most voted for one, so that a single path descending to a
+    // distant patch of ground cannot move the base of the tree
+    const double max_root_radius = 2.0 * pixel_width;
+    for (int label = 0; label < num_labels; label++)
+    {
+      if (best_root[label] == -1)
+      {
+        continue;
+      }
+      std::vector<int> close_roots;
+      for (auto &root : roots_set[label])
+      {
+        Eigen::Vector3d dif = points[root].pos - points[best_root[label]].pos;
+        dif[2] = 0.0;
+        if (dif.norm() <= max_root_radius)
+        {
+          close_roots.push_back(root);
+        }
+      }
+      roots_set[label] = close_roots;
+    }
+    // a label whose points all descend through another label's root points has no root of its own,
+    // so we use its entry points (the points where its paths cross into the label) instead
+    std::vector<bool> needs_entry_points(num_labels);
+    for (int label = 0; label < num_labels; label++)
+    {
+      needs_entry_points[label] = roots_set[label].empty();
+    }
+    for (int i = 0; i < num_cloud_points; i++)
+    {
+      const int label = labels[i];
+      if (label == -1 || !needs_entry_points[label])
+      {
+        continue;
+      }
+      const int parent = points[i].parent;
+      if (parent == -1 || path_labels[parent] != label)
+      {
+        roots_set[label].push_back(i);
+      }
+    }
+    // lastly remove the labels that are too small to count as a tree
+    std::vector<std::vector<int>> result;
+    int num_short = 0;
+    for (int label = 0; label < num_labels; label++)
+    {
+      if (roots_set[label].empty())
+      {
+        continue;
+      }
+      if (max_heights[label] < height_min)
+      {
+        num_short++;
+        continue;
+      }
+      result.push_back(roots_set[label]);
+    }
+    if (num_short > 0)
+    {
+      std::cout << "removed " << num_short << " labelled trees shorter than the " << height_min
+                << " m minimum height" << std::endl;
+    }
+    return result;
+  }
 
   // next we want to segment the paths into separate trees. To do this we find the number of points and
   // the maximum height of points that come from each cell index

raylib/extraction/raysegment.h

@@ -42,9 +42,13 @@ struct RAYLIB_EXPORT Vertex
 /// @c max_diameter maximum diameter of a tree trunk
 /// @c distance_limit maximum distance between points that can be connected
 /// @c gravity_factor controls how far laterally the shortest paths can travel
+/// @c point_labels optional tree label per ray in @c cloud. When supplied, one tree is
+/// generated per distinct label, rather than segmenting the trees from the point distribution. Negative
+/// labels mean unlabelled, these points are attached to whichever labelled tree reaches them
 std::vector<std::vector<int>> RAYLIB_EXPORT getRootsAndSegment(std::vector<Vertex> &points, const Cloud &cloud, const Mesh &mesh,
                                                                double max_diameter, double distance_limit, double height_min,
-                                                               double gravity_factor);
+                                                               double gravity_factor,
+                                                               const std::vector<int> *point_labels = nullptr);
 
 }  // namespace ray
 #endif  // RAYLIB_RAYSEGMENT_H

raylib/extraction/raytrees.cpp

@@ -28,13 +28,16 @@ TreesParams::TreesParams()
 /// The main reconstruction algorithm
 /// It is based on finding the shortest paths using Djikstra's algorithm, followed
 /// by an agglomeration of paths, with repeated splitting from root to tips
-Trees::Trees(Cloud &cloud, const Eigen::Vector3d &offset, const Mesh &mesh, const TreesParams &params, bool verbose)
+Trees::Trees(Cloud &cloud, const Eigen::Vector3d &offset, const Mesh &mesh, const TreesParams &params, bool verbose,
+             const std::vector<int> *point_labels)
 {
   // firstly, get the full set of shortest paths from ground to tips, and the set of roots
   params_ = &params;
+  point_labels_ = point_labels;
 
-  std::vector<std::vector<int>> roots_list = getRootsAndSegment(
-    points_, cloud, mesh, params_->max_diameter, params_->distance_limit, params_->height_min, params_->gravity_factor);
+  std::vector<std::vector<int>> roots_list =
+    getRootsAndSegment(points_, cloud, mesh, params_->max_diameter, params_->distance_limit, params_->height_min,
+                       params_->gravity_factor, point_labels_);
 
   // Now we want to convert these paths into a set of branch sections, from root to tips
   // splitting as we go up...
@@ -145,7 +148,8 @@ Trees::Trees(Cloud &cloud, const Eigen::Vector3d &offset, const Mesh &mesh, cons
     sections_[sec_].tip = best_tip;
     sections_[sec_].ends = best_ends;
     nodes = best_nodes;
-    if (sections_[sec_].split_count < 2)
+    // when the trees are pre-labelled we trust the labels, so we don't split a trunk into several trees
+    if (!point_labels_ && sections_[sec_].split_count < 2)
     {
       double thickness = best_dist; 
       bool points_removed = false;

raylib/extraction/raytrees.h

@@ -44,8 +44,11 @@ class RAYLIB_EXPORT Trees
 {
 public:
   /// Constructs the piecewise cylindrical tree structures from the input ray cloud @c cloud
-  /// The ground @c mesh defines the ground and @params are used to control the reconstruction
-  Trees(Cloud &cloud, const Eigen::Vector3d &offset, const Mesh &mesh, const TreesParams &params, bool verbose);
+  /// The ground @c mesh defines the ground and @params are used to control the reconstruction.
+  /// When @c point_labels is supplied (one tree label per ray in @c cloud, -1 for unlabelled) the
+  /// cloud is treated as pre-segmented, so only the tree models are estimated, one per label
+  Trees(Cloud &cloud, const Eigen::Vector3d &offset, const Mesh &mesh, const TreesParams &params, bool verbose,
+        const std::vector<int> *point_labels = nullptr);
 
   /// save the trees representation to a text file
   bool save(const std::string &filename, const Eigen::Vector3d &offset, bool verbose) const;
@@ -103,6 +106,7 @@ class RAYLIB_EXPORT Trees
   // cached data that is used throughout the processing method
   int sec_;
   const TreesParams *params_;
+  const std::vector<int> *point_labels_; // per-ray tree labels when the cloud is pre-segmented, otherwise null
   std::vector<Vertex> points_;
   double forest_taper_{0};
   double forest_weight_{0};