ArVizDrawer.cpp

Go to the documentation of this file.

#include "ArVizDrawer.h"
 
#include <Eigen/Core>
#include <Eigen/Geometry>
 
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
 
#include <SimoxUtility/iterator/xy_index_range_iterator.h>
#include <SimoxUtility/shapes/OrientedBox.h>
 
#include <ArmarXCore/core/logging/Logging.h>
 
#include "RobotAPI/components/ArViz/Client/elements/PointCloud.h"
#include <RobotAPI/components/ArViz/Client/Client.h>
#include <RobotAPI/components/ArViz/Client/Elements.h>
 
namespace armarx::laser_scanner_simulation
{
 
    template <typename EigenVectorT>
    EigenVectorT
    toCartesian(const LaserScanStep& laserScanStep)
    {
        EigenVectorT point = EigenVectorT::Zero();
 
        point.x() = laserScanStep.distance * std::cos(laserScanStep.angle);
        point.y() = laserScanStep.distance * std::sin(laserScanStep.angle);
 
        return point;
    }
 
    void
    ArVizDrawer::drawBoundingBoxes(const std::vector<VirtualRobot::BoundingBox>& boundingBoxes)
    {
 
        auto layer = this->layer("bounding_boxes");
        int i = 0;
 
        for (const VirtualRobot::BoundingBox& bb : boundingBoxes)
        {
            const Eigen::Vector3f center((bb.getMin() + bb.getMax()) / 2);
            const Eigen::Vector3f extents = (bb.getMax() - bb.getMin());
 
            const auto box = viz::Box("bb_" + std::to_string(i++))
                                 .position(center)
                                 .size(extents)
                                 .color(0, 255, 0, 255);
            layer.add(box);
        }
 
        commit(layer);
    }
 
    class ColoredMeshGrid
    {
    public:
        using Color = Eigen::Vector4i;
        using Position = Eigen::Vector3f;
 
        using ColorGrid = std::vector<std::vector<Color>>;
        using VertexGrid = std::vector<std::vector<Position>>;
 
        ColoredMeshGrid(Eigen::Array2i numCells) :
            numCells(numCells),
            vertices(numCells.x(), std::vector<Position>(numCells.y(), Position::Zero())),
            colors(numCells.x(), std::vector<Color>(numCells.y(), Color::Zero()))
        {
        }
 
        void
        setColoredVertex(const Eigen::Array2i& idx, const Position& position, const Color& color)
        {
            vertices.at(idx.x()).at(idx.y()) = position;
            colors.at(idx.x()).at(idx.y()) = color;
        }
 
        void
        apply(const auto& posFn, const auto& colorFn)
        {
            for (int x = 0; x < numCells.x(); x++)
            {
                for (int y = 0; y < numCells.y(); y++)
                {
                    const Eigen::Array2i idx{x, y};
                    setColoredVertex(idx, posFn(idx), colorFn(idx));
                }
            }
        }
 
        std::vector<std::vector<viz::data::Color>>
        getColors()
        {
            std::vector<std::vector<viz::data::Color>> colors(
                numCells.x(), std::vector<viz::data::Color>(numCells.y(), viz::Color::black()));
 
            auto toVizColor = [](const ColoredMeshGrid::Color& c)
            { return viz::Color(c.x(), c.y(), c.z(), c.w()); };
 
            auto xyIter = simox::iterator::XYIndexRangeIterator({numCells.x(), numCells.y()});
            std::for_each(xyIter.begin(),
                          xyIter.end(),
                          [&](const auto& p)
                          {
                              const auto& [x, y] = p;
                              colors[x][y] = toVizColor(this->colors[x][y]);
                          });
 
            return colors;
        }
 
        const VertexGrid&
        getVertices() const
        {
            return vertices;
        }
 
    private:
        Eigen::Array2i numCells;
 
        VertexGrid vertices;
        ColorGrid colors;
    };
 
    void
    ArVizDrawer::drawOccupancyGrid(const OccupancyGrid& grid, const float boxPosZ)
    {
        ARMARX_INFO << "Occupancy grid size " << Eigen::Array2i{grid.sizeX, grid.sizeY};
 
        auto layer = this->layer("occupancy_grid");
 
        ColoredMeshGrid meshGrid(Eigen::Array2i{grid.sizeX, grid.sizeY});
 
        const auto posFn = [&](const Eigen::Array2i& idx) -> ColoredMeshGrid::Position
        {
            const auto pos2d = grid.getCentralPosition(idx.x(), idx.y());
            return {pos2d.x(), pos2d.y(), 0.F};
        };
 
        const ColoredMeshGrid::Color colorBlack{0, 0, 0, 255};
        const ColoredMeshGrid::Color colorWhite{255, 255, 255, 255};
 
        const auto colorFn = [&](const Eigen::Array2i& idx) -> ColoredMeshGrid::Color
        {
            if (grid.isOccupied(static_cast<std::size_t>(idx.x()),
                                static_cast<std::size_t>(idx.y())))
            {
                ARMARX_IMPORTANT << "Cell is occupied";
            }
            return grid.isOccupied(static_cast<std::size_t>(idx.x()),
                                   static_cast<std::size_t>(idx.y()))
                       ? colorBlack
                       : colorWhite;
        };
 
        meshGrid.apply(posFn, colorFn);
 
        const auto mesh = viz::Mesh("grid")
                              .position(Eigen::Vector3f{0.F, 0.F, boxPosZ})
                              .grid2D(meshGrid.getVertices(), meshGrid.getColors());
 
        layer.add(mesh);
        commit(layer);
    }
 
    pcl::PointCloud<pcl::PointXYZ>
    toPointCloud(const LaserScan& scan)
    {
        pcl::PointCloud<pcl::PointXYZ> cloud;
        cloud.points.reserve(scan.size());
 
        std::transform(scan.begin(),
                       scan.end(),
                       std::back_inserter(cloud.points),
                       [&](const auto& scanStep) -> pcl::PointXYZ
                       {
                           const auto rangePoint = toCartesian<Eigen::Vector3f>(scanStep); // [mm]
 
                           pcl::PointXYZ pt;
                           pt.x = rangePoint.x();
                           pt.y = rangePoint.y();
                           pt.z = rangePoint.z();
 
                           return pt;
                       });
 
        cloud.width = scan.size();
        cloud.height = 1;
        cloud.is_dense = true;
 
        return cloud;
    }
 
    void
    ArVizDrawer::prepareScan(const LaserScan& scan,
                             const std::string& frame,
                             const Eigen::Affine3f& scannerPose)
    {
 
        const auto pointCloud = toPointCloud(scan);
 
        frames.insert(frame);
 
        // skip the first two colors as they are 'black' and 'white'
        const int id = std::distance(std::begin(frames), frames.find(frame)) + 2;
 
        clouds.emplace(
            frame,
            viz::PointCloud(frame)
                .pose(scannerPose)
                .pointSizeInPixels(5)
                .pointCloud(pointCloud.points, viz::Color(simox::color::GlasbeyLUT::at(id))));
    }
 
    void
    ArVizDrawer::drawScans()
    {
        auto layer = this->layer("scans");
 
        for (const auto& [_, pc] : clouds)
        {
            layer.add(pc);
        }
 
        commit(layer);
 
        // remove all
        clouds.clear();
    }
 
} // namespace armarx::laser_scanner_simulation