src/kinect/segmenter.cpp

Go to the documentation of this file.

#include "ed/kinect/segmenter.h"
 
#include <geolib/sensors/DepthCamera.h>
#include <geolib/Shape.h>
 
#include <rgbd/image.h>
#include <rgbd/view.h>
 
#include <ed/world_model.h>
#include <ed/entity.h>
 
#include <ros/console.h>
 
// Clustering
#include <queue>
#include <ed/convex_hull_calc.h>
#include <sys/resource.h>
 
#include <opencv2/ml.hpp>
#include <pcl/point_types.h>
#include <pcl/point_cloud.h>
#include <pcl/segmentation/extract_clusters.h>
 
#include <Eigen/Dense>
 
#include <ed_sensor_integration/kinect/segmodules/sam_seg_module.h>
#include <bmm/bayesian_mixture_model.hpp>
#include <chrono>
 
// ----------------------------------------------------------------------------------------------------
 
Segmenter::Segmenter(tue::Configuration config)
    : config_(config)
{
    if (config_.readArray("surface_label_map"))
    {
        while (config_.nextArrayItem())
        {
            std::string entity_name, yolo_label;
            if (config_.value("entity", entity_name) && config_.value("yolo_label", yolo_label))
                surface_label_map_[entity_name] = yolo_label;
        }
        config_.endArray();
    }
 
    ROS_INFO("Starting async initialization of segmentation pipeline (YOLO & SAM)...");
    // Note: Atomic flag does not prevent 2 threads from initializing the pipeline simultaneously.
    // This is not needed as the pipeline should only be initialized once.
    init_thread_ = std::thread([this, config]() mutable {
        this->sam_pipeline_.initialize(config);
        ROS_INFO("Async initialization of segmentation pipeline is complete!");
    });
}
 
// ----------------------------------------------------------------------------------------------------
 
Segmenter::~Segmenter()
{
    if (init_thread_.joinable())
    {
        init_thread_.join();
    }
}
 
// ----------------------------------------------------------------------------------------------------
 
namespace
{
// Internal constants (tuning thresholds)
constexpr std::size_t MIN_FILTERED_POINTS = 10;
constexpr double      MIN_RETENTION_RATIO = 0.10;  // 10%
constexpr std::size_t MIN_CLUSTER_POINTS = 100;
 
class DepthRenderer : public geo::RenderResult
{
 
public:
 
    DepthRenderer(cv::Mat& z_buffer_)
        : geo::RenderResult(z_buffer_.cols, z_buffer_.rows), z_buffer(z_buffer_)
    {
    }
 
    void renderPixel(int x, int y, float depth, int /*i_triangle*/)
    {
        float& old_depth = z_buffer.at<float>(y, x);
        if (old_depth == 0 || depth < old_depth)
        {
            old_depth = depth;
        }
    }
 
    cv::Mat& z_buffer;
};
 
}
 
// ----------------------------------------------------------------------------------------------------
 
void Segmenter::removeBackground(cv::Mat& depth_image, const ed::WorldModel& world, const geo::DepthCamera& cam,
                                 const geo::Pose3D& sensor_pose, double background_padding)
{
    // - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    // Render the world model as seen by the depth sensor
 
    cv::Mat depth_model(depth_image.rows, depth_image.cols, CV_32FC1, 0.0);
 
    DepthRenderer res(depth_model);
    for(ed::WorldModel::const_iterator it = world.begin(); it != world.end(); ++it)
    {
        const ed::EntityConstPtr& e = *it;
        if (!e->visual() || !e->has_pose())
            continue;
 
        geo::RenderOptions opt;
        opt.setMesh(e->visual()->getMesh(), sensor_pose.inverse() * e->pose());
        cam.render(opt, res);
    }
 
    // - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    // Filter all points that can be associated with the rendered depth image
 
    unsigned int size = depth_image.rows * depth_image.cols;
    for(unsigned int i = 0; i < size; ++i)
    {
        float& ds = depth_image.at<float>(i);
        float dm = depth_model.at<float>(i);
        if (dm > 0 && ds > 0 && ds > dm - background_padding)
            ds = 0;
    }
}
 
// ----------------------------------------------------------------------------------------------------
 
class MinMaxRenderer : public geo::RenderResult
{
 
public:
 
    MinMaxRenderer(int width, int height) : geo::RenderResult(width, height)
    {
        min_buffer = cv::Mat(height, width, CV_32FC1, 0.0);
        max_buffer = cv::Mat(height, width, CV_32FC1, 0.0);
    }
 
    void renderPixel(int x, int y, float depth, int /*i_triangle*/)
    {
        // TODO: now the renderer can only deal with convex meshes, which means
        // that at each pixel there can only be one minimum and one maximum pixel
        // There is an easy solution for concave meshes: determine which side
        // the triangle points (away or to the camera) and test the pixel in the depth
        // image to be on the correct side. ... etc ...
 
        float& d_min = min_buffer.at<float>(y, x);
        float& d_max = max_buffer.at<float>(y, x);
 
        if (d_min == 0 || depth < d_min)
            d_min = depth;
 
        d_max = std::max(d_max, depth);
    }
 
    cv::Mat min_buffer;
    cv::Mat max_buffer;
 
};
 
// ----------------------------------------------------------------------------------------------------
 
void Segmenter::calculatePointsWithin(const rgbd::Image& image, const geo::Shape& shape,
                                      const geo::Pose3D& shape_pose, cv::Mat& filtered_depth_image) const
{
    // - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    // Render shape
 
    const cv::Mat& depth_image = image.getDepthImage();
 
    rgbd::View view(image, depth_image.cols);
    const geo::DepthCamera& cam_model = view.getRasterizer();
 
    MinMaxRenderer res(depth_image.cols, depth_image.rows);
 
    geo::RenderOptions opt;
    opt.setBackFaceCulling(false);
    opt.setMesh(shape.getMesh(), shape_pose);
 
    cam_model.render(opt, res);
 
    filtered_depth_image = cv::Mat(depth_image.rows, depth_image.cols, CV_32FC1, 0.0);
 
    for(int i = 0; i < depth_image.cols * depth_image.rows; ++i)
    {
        float d = depth_image.at<float>(i);
        if (d <= 0)
            continue;
 
        float d_min = res.min_buffer.at<float>(i);
        float d_max = res.max_buffer.at<float>(i);
 
        if (d_min > 0 && d_max > 0 && d >= d_min && d <= d_max)
            filtered_depth_image.at<float>(i) = d;
    }
}
 
// ----------------------------------------------------------------------------------------------------
 
cv::Mat Segmenter::preprocessRGBForSegmentation(const cv::Mat& rgb_image,
                                                const cv::Mat& filtered_depth_image) const
{
    cv::Mat masked_rgb = cv::Mat::zeros(rgb_image.size(), rgb_image.type());
    for (int y = 0; y < rgb_image.rows; ++y) {
        const float* depth_row = filtered_depth_image.ptr<float>(y); // fast
        cv::Vec3b* out_row = masked_rgb.ptr<cv::Vec3b>(y);
        const cv::Vec3b* rgb_row = rgb_image.ptr<cv::Vec3b>(y);
        for (int x = 0; x < rgb_image.cols; ++x) {
            if (depth_row[x] > 0.f) {
                out_row[x] = rgb_row[x];
            }
        }
    }
    return masked_rgb;
}
// ----------------------------------------------------------------------------------------------------
 
SegmentationResult Segmenter::cluster(const cv::Mat& depth_image, const geo::DepthCamera& cam_model,
                        const geo::Pose3D& sensor_pose, std::vector<EntityUpdate>& clusters, const cv::Mat& rgb_image,
                        const std::string& area_description, bool verbose)
{
    int width = depth_image.cols;
    int height = depth_image.rows;
    ROS_DEBUG("Cluster with depth image of size %i, %i", width, height);
 
    // Extract area name and entity from area_description (e.g. "on_top_of" and "dinner_table" from "on_top_of dinner_table")
    std::string area_name;
    std::string area_entity;
    std::string ignore_label = "";
 
    // this code block is an example for manually ignoring the supporting surface when looking for objects on top of it.
 
    //index of the space that separates the 2 words in the area description (e.g. "on_top_of dinner_table")
    // const std::size_t i_space = area_description.find(' ');
    // if (i_space != std::string::npos)
    // {
    //     area_name   = area_description.substr(0, i_space);
    //     area_entity = area_description.substr(i_space + 1);
 
    //     if (area_name == "on_top_of")
    //     {
    //         const auto it = surface_label_map_.find(area_entity);
    //         if (it != surface_label_map_.end())
    //         {
    //             ignore_label = it->second;
    //         }
    //     }
    // }
 
    // Pass verbose flag to measure latency in the segmentation pipeline
    SegmentationResult seg_result = sam_pipeline_.process(rgb_image.clone(), depth_image, ignore_label, verbose);
    std::vector<cv::Mat>& masks = seg_result.masks;
 
    ROS_DEBUG("Creating clusters");
 
    // Pre-allocate temporary storage (one per mask, avoid push_back races)
    std::vector<EntityUpdate> temp_clusters(masks.size());
    std::vector<bool> valid_cluster(masks.size(), false);
 
    // BMM point cloud denoising
    GMMParams params;
    if (config_.readGroup("bmm"))
    {
        config_.value("psi0", params.psi0, tue::config::OPTIONAL);
        config_.value("nu0", params.nu0, tue::config::OPTIONAL);
        config_.value("alpha", params.alpha, tue::config::OPTIONAL);
        config_.value("kappa0", params.kappa0, tue::config::OPTIONAL);
        config_.endGroup();
    }
    // BMM timing: per-mask latencies stored for aggregation (thread-safe via indexing)
    std::vector<double> bmm_latencies_ms(masks.size(), 0.0);
 
    // Parallel loop - each iteration is independent
    #pragma omp parallel for schedule(dynamic)
    for (size_t i = 0; i < masks.size(); ++i)
    {
        //Resize mask if needed so for the image to be the same size as the depth image (in case SAM produces a different size mask)
        const cv::Mat& mask_orig = masks[i];
 
        // Empty mask: SAM failed for this box (placeholder pushed to preserve alignment)
        if (mask_orig.empty())
        {
            ROS_DEBUG("Skipping cluster %zu: SAM returned empty mask (segmentation failure for this box)", i);
            continue;
        }
        // When looking on_top_of a surface, skip the supporting surface itself so it is
        // never extracted as a cluster or passed through BMM.  The entity-to-YOLO-label
        // mapping is read from the "surface_label_map" config array.
        if (area_name == "on_top_of" && i < seg_result.labels.size())
        {
            const auto it = surface_label_map_.find(area_entity);
            if (it != surface_label_map_.end() && seg_result.labels[i] == it->second)
            {
                ROS_WARN("Skipping cluster %zu: label '%s' is the supporting surface for area '%s' and area description '%s'",
                          i, seg_result.labels[i].c_str(), area_name.c_str(), area_description.c_str());
                continue;
            }
        }
 
        cv::Mat mask;
        if (mask_orig.rows != height || mask_orig.cols != width)
            cv::resize(mask_orig, mask, cv::Size(width, height), 0, 0, cv::INTER_NEAREST);
        else
            mask = mask_orig;
 
        EntityUpdate cluster;  // local to this thread
 
        // Extract points from mask
        for(int y = 0; y < mask.rows; ++y) {
            for(int x = 0; x < mask.cols; ++x) {
                if (mask.at<unsigned char>(y, x) > 0) {
                    unsigned int pixel_idx = y * width + x;
                    float d = depth_image.at<float>(pixel_idx);
 
                    if (d > 0 && std::isfinite(d)) {
                        cluster.pixel_indices.push_back(pixel_idx);
                        cluster.points.push_back(cam_model.project2Dto3D(x, y) * d);
                    }
                }
            }
        }
 
        // Skip small clusters (< 100 points based on MIN_CLUSTER_POINTS)
        if (cluster.pixel_indices.size() < MIN_CLUSTER_POINTS) {
            if (verbose)
            {
                const std::string label = (i < seg_result.labels.size()) ? seg_result.labels[i] : "?";
                ROS_WARN("We reject cluster %zu with label '%s' because it has only %zu points", i, label.c_str(), cluster.pixel_indices.size());
            }
            continue;  // valid_cluster[i] remains false
        }
 
        if (verbose)
        {
            const std::string label = (i < seg_result.labels.size()) ? seg_result.labels[i] : "?";
            ROS_WARN("Cluster %zu with label '%s': %zu points", i, label.c_str(), cluster.points.size());
        }
 
        // BMM inference with timing
        auto bmm_start = std::chrono::high_resolution_clock::now();
 
        MAPGMM gmm(2, cluster.points, params); // 2 components: object + outliers
        gmm.fit(cluster.points, sensor_pose);
 
        auto bmm_end = std::chrono::high_resolution_clock::now();
        bmm_latencies_ms[i] = std::chrono::duration<double, std::milli>(bmm_end - bmm_start).count();
 
        // Get component assignments and inlier component
        std::vector<int> labels = gmm.get_labels();
        int inlier_component = gmm.get_inlier_component();
 
        // Filter points
        std::vector<geo::Vec3> filtered_points;
        std::vector<geo::Vec3> outlier_points;  // Only populate if needed
 
        for (size_t j = 0; j < labels.size(); j++)
        {
            if (labels[j] == inlier_component)
            {
                filtered_points.push_back(cluster.points[j]);
            }
            else if (verbose)
            {
                outlier_points.push_back(cluster.points[j]);
            }
        }
 
        // Safety check: only use filtered points if we retained enough
        if (filtered_points.size() > MIN_FILTERED_POINTS &&
            filtered_points.size() > MIN_RETENTION_RATIO * cluster.points.size()) {
            // Use filtered points
            cluster.points = filtered_points;
            if (verbose)
            {
                cluster.outlier_points = outlier_points;
                // Transform outlier points to map frame
                // for (size_t j = 0; j < cluster.outlier_points.size(); ++j) {
                //     cluster.outlier_points[j] = sensor_pose * cluster.outlier_points[j];
                // }
            }
        }
        else
        {
            // Safety check failed - keep original unfiltered points
            // Don't populate outlier_points since we're not using the GMM result
            ROS_DEBUG("GMM filtering rejected: retained %zu/%zu points",
                      filtered_points.size(), cluster.points.size());
        }
 
        // Calculate convex hull
        float z_min = 1e9;
        float z_max = -1e9;
        std::vector<geo::Vec2f> points_2d(cluster.points.size());
 
        for(unsigned int j = 0; j < cluster.points.size(); ++j)
        {
            const geo::Vec3& p = cluster.points[j];
 
            // Transform sensor point to map frame
            geo::Vector3 p_map = sensor_pose * p;
            points_2d[j] = geo::Vec2f(p_map.x, p_map.y);
            z_min = std::min<float>(z_min, p_map.z);
            z_max = std::max<float>(z_max, p_map.z);
        }
 
        ed::convex_hull::create(points_2d, z_min, z_max, cluster.chull, cluster.pose_map);
        cluster.chull.complete = false;
 
        // Assign YOLO classification label to this cluster
        if (i < seg_result.labels.size())
        {
            cluster.label = seg_result.labels[i];
            cluster.classification_confidence = seg_result.confidences[i];
            if (verbose)
            {
                ROS_INFO("Cluster %zu classified as '%s' with confidence %.2f", i, cluster.label.c_str(), cluster.classification_confidence);
            }
        }
 
        // Store in thread-safe pre-allocated array
        temp_clusters[i] = cluster;
        valid_cluster[i] = true;
    }
 
    // Sequential section: collect valid clusters
    clusters.clear();
    clusters.reserve(masks.size());
    for (size_t i = 0; i < temp_clusters.size(); ++i) {
        if (valid_cluster[i]) {
            clusters.push_back(std::move(temp_clusters[i]));
        }
    }
 
    // Log BMM latency statistics if verbose is enabled
    if (verbose)
    {
        double bmm_total_ms = 0.0;
        double bmm_max_ms = 0.0;
        int bmm_count = 0;
        for (size_t i = 0; i < bmm_latencies_ms.size(); ++i)
        {
            if (bmm_latencies_ms[i] > 0.0)
            {
                bmm_total_ms += bmm_latencies_ms[i];
                bmm_max_ms = std::max(bmm_max_ms, bmm_latencies_ms[i]);
                bmm_count++;
            }
        }
        double bmm_avg_ms = (bmm_count > 0) ? (bmm_total_ms / bmm_count) : 0.0;
 
        ROS_WARN("\n=== BMM Latency (parallel execution) ===\n");
        ROS_WARN("  BMM fits:      %d masks\n", bmm_count);
        ROS_WARN("  BMM avg:       %7.2f ms/mask\n", bmm_avg_ms);
        ROS_WARN("  BMM max:       %7.2f ms (slowest mask)\n", bmm_max_ms);
        ROS_WARN("  BMM total:     %7.2f ms (sequential equivalent)\n", bmm_total_ms);
    }
 
    return seg_result;
}