cspace.cpp

Go to the documentation of this file.

#include <se2ez/core/log.h>
#include <se2ez/core/math.h>
#include <se2ez/core/joint.h>
#include <se2ez/core/frame.h>
#include <se2ez/core/geometry.h>
#include <se2ez/core/state.h>
#include <se2ez/core/box2d/collision.h>
#include <se2ez/core/colormap.h>

#include <se2ez/core/cspace.h>

using namespace se2ez;

///
/// CSpaceGrid
///

const unsigned int CSpaceGrid::cores_{std::thread::hardware_concurrency()};

CSpaceGrid::CSpaceGrid(const RobotPtr &robot, unsigned int jointA, unsigned int jointB, unsigned int width,
                       unsigned int height, Mode mode)
  : robot_(robot)
  , indexA_(jointA)
  , indexB_(jointB)
  , limitA_({robot->getLowerLimits()[indexA_], robot->getUpperLimits()[indexA_]})
  , limitB_({robot->getLowerLimits()[indexB_], robot->getUpperLimits()[indexB_]})
  , contA_(robot->getContinuity()[indexA_])
  , contB_(robot->getContinuity()[indexB_])
  , width_(width)
  , height_(height)
  , mode_(mode)
  , image_(std::make_shared<ompl::PPM>())
  , state_(std::make_shared<State>(robot))
  , temp_(std::make_shared<State>(robot))
  , newState_(cores_, false)
  , working_(cores_, false)
  , threads_(cores_)
{
    initialize();
}

CSpaceGrid::CSpaceGrid(const RobotPtr &robot, const std::string &jointA, const std::string &jointB,
                       unsigned int width, unsigned int height, Mode mode)
  : CSpaceGrid(robot, robot->getFrameDataConst(jointA)->index, robot->getFrameDataConst(jointB)->index,  //
               width, height, mode)
{
}

void CSpaceGrid::initialize()
{
    // Setup PPM
    image_->setWidth(width_);
    image_->setHeight(height_);
    image_->getPixels().resize(width_ * height_);

    for (unsigned int c = 0; c < cores_; ++c)
    {
        threads_[c] = new std::thread([&, c]() {
            auto cm = std::make_shared<Box2DCollisionManager>(robot_);
            auto state = std::make_shared<State>(robot_);
            double &jA = state->data[indexA_];
            double &jB = state->data[indexB_];

            while (true)
            {
                Chunk chunk;
                GridCallback callback = {};

                // Pop a new chunk off
                {
                    // Wait until there are chunks ready or it's time to exit.
                    std::unique_lock<std::mutex> slock(lock_);
                    work_.wait(slock, [&] { return !chunks_.empty() || done_; });

                    // All done, time to exit.
                    if (done_)
                        break;

                    // Copy state contents and mark done if needed.
                    {
                        std::unique_lock<std::mutex> slock(statelock_);
                        if (newState_[c])
                        {
                            newState_[c] = false;
                            state->copy(state_);
                        }
                    }

                    callback = callback_;

                    // Copy chunk
                    chunk = chunks_.front();
                    chunks_.pop();

                    // SE2EZ_DEBUG("Thread %5%: Popping image chunk %1%x%2% -> %3%x%4%",  //
                    //             chunk.x, chunk.y, chunk.x + chunk.w, chunk.y + chunk.h, c);
                }

                working_[c] = true;

                // Iterate through pixels that need to be filled.
                for (unsigned int i = chunk.x; i < chunk.x + chunk.w; ++i)
                {
                    jA = math::remap(0, width_, i, limitA_[0], limitA_[1]);

                    for (unsigned int j = chunk.y; j < chunk.y + chunk.h; ++j)
                    {
                        jB = math::remap(0, height_, j, limitB_[0], limitB_[1]);

                        state->dirty = true;
                        auto &pixel = image_->getPixel(i, j);

                        bool collision = false;
                        unsigned char rc = 0, gc = 0, bc = 0;
                        switch (mode_)
                        {
                            case DISTANCE:
                            case DISTANCE_COLOR:
                            {
                                double r, g, b;
                                auto sd = cm->distance(state);
                                double d = sd.distance;

                                if (d < 0)
                                {
                                    collision = true;
                                    if (d < minDepth_)
                                        d = minDepth_;

                                    if (mode_ == DISTANCE_COLOR)
                                    {
                                        grayscale(math::remap(0, minDepth_, d, 0.25, 1), r, g, b);
                                        auto color = getFrameColor(sd.collider);

                                        r = color[0] * r;
                                        g = color[1] * g;
                                        b = color[2] * b;
                                    }
                                    else
                                        grayscale(math::remap(0, minDepth_, d, 0, 1), r, g, b);
                                }
                                else
                                {
                                    if (d > maxDepth_)
                                        d = maxDepth_;
                                    turbo(math::remap(0, maxDepth_, d, 0, 1), r, g, b);
                                }

                                rc = std::round(255 * r);
                                gc = std::round(255 * g);
                                bc = std::round(255 * b);

                                break;
                            }
                            case COLOR:
                            {
                                auto sd = cm->distance(state);
                                if (sd.distance >= 0)
                                    rc = gc = bc = 255;
                                else
                                {
                                    collision = true;
                                    auto color = getFrameColor(sd.collider);

                                    rc = std::round(255 * color[0]);
                                    gc = std::round(255 * color[1]);
                                    bc = std::round(255 * color[2]);
                                }

                                break;
                            }
                            case COLLISION:
                                collision = cm->collide(state);
                                rc = gc = bc = (collision) ? 0 : 255;
                                break;
                        }

                        pixel = ompl::PPM::Color{rc, gc, bc};

                        if (callback)
                            callback(pixel, state, collision);
                    }
                }

                working_[c] = false;

                // Notify that we might be complete.
                if (chunks_.empty())
                    complete_.notify_one();
            }
        });
    }
}

CSpaceGrid::~CSpaceGrid()
{
    // Let threads know we are done.
    done_ = true;
    work_.notify_all();

    for (unsigned int c = 0; c < cores_; ++c)
    {
        threads_[c]->join();
        delete threads_[c];
    }
}

void CSpaceGrid::emplaceChunks()
{
    cancel();

    std::unique_lock<std::mutex> slock(lock_);

    // Pump some chunks.
    for (unsigned int x = 0; x < width_; x += xChunk_)
        for (unsigned int y = 0; y < height_; y += yChunk_)
        {
            unsigned int remainX = std::min<unsigned int>({xChunk_, width_ - x});
            unsigned int remainY = std::min<unsigned int>({yChunk_, height_ - y});
            // SE2EZ_DEBUG("Emplacing image chunk %1%x%2% -> %3%x%4%", x, y, x + remainX, y + remainY);
            chunks_.emplace(x, y, remainX, remainY);
        }

    maxChunks_ = chunks_.size();

    slock.unlock();
    work_.notify_all();
}

void CSpaceGrid::compute(bool block, const StatePtr &base, const GridCallback &callback)
{
    // Notify threads there is a new state to copy.
    if (base)
    {
        std::unique_lock<std::mutex> slock(statelock_);
        state_->copy(base);

        for (unsigned int i = 0; i < cores_; ++i)
            newState_[i] = true;
    }

    callback_ = {};
    if (callback)
        callback_ = callback;

    // Reboot work queue.
    emplaceChunks();

    // Possibly block until exit.
    if (block)
        waitToComplete();
}

void CSpaceGrid::waitToComplete()
{
    // SE2EZ_DEBUG("Waiting for CSpace Grid computation...");

    std::unique_lock<std::mutex> slock(lock_);

    // Wait until all threads are done working.
    complete_.wait(slock, [&] { return not isComputing(); });

    // SE2EZ_DEBUG("Done waiting for CSpace Grid computation!");
}

ompl::PPM::Color &CSpaceGrid::getPixel(const StatePtr &state)
{
    return getPixel(state->data[indexA_], state->data[indexB_]);
}

CSpaceGrid::GridCoords CSpaceGrid::getCoords(const StatePtr &state)
{
    return getCoords(state->data[indexA_], state->data[indexB_]);
}

ompl::PPM::Color &CSpaceGrid::getPixel(double i, double j)
{
    const auto &coord = getCoords(i, j);
    return image_->getPixel(coord.x, coord.y);
}

CSpaceGrid::GridCoords CSpaceGrid::getCoords(double i, double j)
{
    double x = math::remap(limitA_[0], limitA_[1], i, 0, width_);
    double y = math::remap(limitB_[0], limitB_[1], j, 0, height_);

    return {.x = (unsigned int)x, .y = (unsigned int)y};
}

CSpaceGrid::RobotCoords CSpaceGrid::getCoords(unsigned int i, unsigned int j)
{
    double x = math::remap(0, width_, i, limitA_[0], limitA_[1]);
    double y = math::remap(0, height_, j, limitB_[0], limitB_[1]);

    return {.x = x, .y = y};
}

std::vector<CSpaceGrid::GridLine> CSpaceGrid::getLines(const StatePtr &a, const StatePtr &b)
{
    std::vector<GridLine> lines;
    lines.push_back({.a = getCoords(a), .b = getCoords(b)});

    auto scratch = std::make_shared<State>(robot_);
    if (contA_)
    {
        std::vector<GridLine> newLines;
        for (const auto &line : lines)
        {
            unsigned int ax = line.a.x, ay = line.a.y, bx = line.b.x, by = line.b.y;
            if (ax > bx)
            {
                std::swap(ax, bx);
                std::swap(ay, by);
            }

            if ((bx - ax) <= width_ / 2)
            {
                newLines.emplace_back(line);
                continue;
            }

            double at = ax;
            double bt = width_ - bx;
            double tt = at + bt;

            unsigned int ma = ay + ((int)by - (int)ay) * at / tt;

            GridCoords an = {.x = 0, .y = ma};
            GridCoords aa = {.x = ax, .y = ay};

            GridCoords bn = {.x = width_, .y = ma};
            GridCoords bb = {.x = bx, .y = by};

            newLines.push_back({.a = an, .b = aa});
            newLines.push_back({.a = bn, .b = bb});
        }
        lines = newLines;
    }

    if (contB_)
    {
        std::vector<GridLine> newLines;
        for (const auto &line : lines)
        {
            unsigned int ax = line.a.x, ay = line.a.y, bx = line.b.x, by = line.b.y;
            if (ay > by)
            {
                std::swap(ax, bx);
                std::swap(ay, by);
            }

            if ((by - ay) <= height_ / 2)
            {
                newLines.emplace_back(line);
                continue;
            }

            double at = ay;
            double bt = height_ - by;
            double tt = at + bt;

            unsigned int ma = ax + ((int)bx - (int)ax) * at / tt;

            GridCoords an = {.x = ma, .y = 0};
            GridCoords aa = {.x = ax, .y = ay};

            GridCoords bn = {.x = ma, .y = height_};
            GridCoords bb = {.x = bx, .y = by};

            newLines.push_back({.a = an, .b = aa});
            newLines.push_back({.a = bn, .b = bb});
        }

        lines = newLines;
    }

    return lines;
}

std::shared_ptr<ompl::PPM> CSpaceGrid::getImage() const
{
    return image_;
}

bool CSpaceGrid::isComputing() const
{
    bool working = false;
    for (unsigned int i = 0; i < cores_; ++i)
        working |= working_[i];
    return working or not chunks_.empty();
}

void CSpaceGrid::cancel()
{
    std::unique_lock<std::mutex> slock(lock_);

    // Clear current work queue.
    while (!chunks_.empty())
        chunks_.pop();
}

double CSpaceGrid::distanceToSlice(const StatePtr &state)
{
    temp_->copy(state);
    temp_->data[indexA_] = 0;
    temp_->data[indexB_] = 0;

    return temp_->data.norm();
}

double CSpaceGrid::progress() const
{
    if (not maxChunks_)
        return 1.;

    return (double)(maxChunks_ - chunks_.size()) / maxChunks_;
}

Eigen::Vector4d CSpaceGrid::getFrameColor(const std::string &frame) const
{
    auto framep = robot_->getFrame(frame);
    Eigen::Vector4d color{0, 0, 0, 0};
    for (const auto &geometry : framep->getGeometry())
        color += geometry->color;

    color /= (double)framep->getGeometry().size();
    return color;
}

unsigned int CSpaceGrid::getWidth() const
{
    return width_;
}

unsigned int CSpaceGrid::getHeight() const
{
    return height_;
}