File switchboard_bindings.hpp

File List > plugins > semantic_python > switchboard_bindings.hpp

Go to the documentation of this file

#pragma once

#include "plugin.hpp"

#include 
#include 
#include 
#include 
#include 
#include 
#include 

namespace ILLIXR {

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

static pybind11::array_t<uint8_t> to_numpy_flat_safe(std::shared_ptr<const data_format::semantic_frame> owner,
                                                     const uint8_t* data, size_t size) {
    pybind11::capsule base(new std::shared_ptr<const data_format::semantic_frame>(std::move(owner)), [](void* p) {
        delete static_cast<std::shared_ptr<const data_format::semantic_frame>*>(p);
    });

    return pybind11::array_t<uint8_t>({static_cast<pybind11::ssize_t>(size)}, {static_cast<pybind11::ssize_t>(1)}, data, base);
}

static pybind11::array_t<float> to_numpy_4x4_safe(std::shared_ptr<const data_format::semantic_frame> owner, const float* data) {
    pybind11::capsule base(new std::shared_ptr<const data_format::semantic_frame>(std::move(owner)), [](void* p) {
        delete static_cast<std::shared_ptr<const data_format::semantic_frame>*>(p);
    });

    return pybind11::array_t<float>(
        {static_cast<pybind11::ssize_t>(4), static_cast<pybind11::ssize_t>(4)},
        {static_cast<pybind11::ssize_t>(sizeof(float) * 4), static_cast<pybind11::ssize_t>(sizeof(float))}, data, base);
}

// Convert a camera_intrinsics struct to a 1D numpy float32 array [fx, fy, cx, cy].
// Copies the four floats by value — the struct is not contiguous in memory, so
// zero-copy is not applicable here.
static pybind11::array_t<float> intrinsics_to_numpy(const data_format::camera_intrinsics& intr) {
    auto result = pybind11::array_t<float>(4);
    auto buf    = result.mutable_unchecked<1>();
    buf(0)      = intr.fx;
    buf(1)      = intr.fy;
    buf(2)      = intr.cx;
    buf(3)      = intr.cy;
    return result;
}

// Wrap a DecodedFrameCache::Entry's RGB vector as a (H, W, 3) uint8 numpy
// array without copying. The capsule keeps the Entry alive via a shared_ptr
// to the cache entry data. Since Entry is owned by the cache (fixed array),
// we share-own a copy of the rgb vector data instead.
static pybind11::array_t<uint8_t> entry_to_numpy(const decode::decoded_frame_cache::Entry* entry) {
    if (entry == nullptr || entry->rgb.empty())
        return pybind11::array_t<uint8_t>();

    // Copy the vector into a heap-allocated buffer owned by the capsule.
    // This is one memcpy but avoids lifetime issues with the cache slot
    // being overwritten while Python holds a reference.
    auto*             buf = new std::vector<uint8_t>(entry->rgb);
    pybind11::capsule base(buf, [](void* p) {
        delete static_cast<std::vector<uint8_t>*>(p);
    });

    return pybind11::array_t<uint8_t>({static_cast<pybind11::ssize_t>(entry->height),
                                       static_cast<pybind11::ssize_t>(entry->width), static_cast<pybind11::ssize_t>(3)},
                                      {static_cast<pybind11::ssize_t>(entry->width * 3), static_cast<pybind11::ssize_t>(3),
                                       static_cast<pybind11::ssize_t>(1)},
                                      buf->data(), base);
}

// ---------------------------------------------------------------------------
// SemanticDataReader proxy
// ---------------------------------------------------------------------------

struct py_semantic_data_reader {
    decode::decoded_frame_cache*     cache;
    decode::semantic_metadata_cache* metadata_cache;

    [[nodiscard]] pybind11::object get() const {
        auto entry_opt = cache ? cache->latest() : std::nullopt;
        if (!entry_opt)
            return pybind11::none();

        auto meta_opt = metadata_cache ? metadata_cache->find(entry_opt->frame_number) : std::nullopt;
        if (!meta_opt)
            return pybind11::none();

        // entry_opt->rgb is our own copy — safe to use without any lock.
        // meta_opt->data is a shared_ptr — refcount keeps the frame alive.
        auto val = meta_opt->data;

        pybind11::dict data;
        data["image"]              = entry_to_numpy(&*entry_opt);
        data["frame_number"]       = entry_opt->frame_number;
        data["image_width"]        = val->intrinsics.width;
        data["image_height"]       = val->intrinsics.height;
        data["depth"]              = to_numpy_flat_safe(val, val->depth.data(), val->depth.size());
        data["depth_width"]        = val->depth_intrinsics.width;
        data["depth_height"]       = val->depth_intrinsics.height;
        data["depth_near_z"]       = val->depth_near_z;
        data["intrinsics"]         = intrinsics_to_numpy(val->intrinsics);
        data["depth_intrinsics"]   = intrinsics_to_numpy(val->depth_intrinsics);
        data["rgb_camera_pose"]    = to_numpy_4x4_safe(val, val->rgb_camera_pose);
        data["depth_pose"]         = to_numpy_4x4_safe(val, val->depth_pose);
        data["max_depth_m"]        = val->max_depth;
        data["rgb_timestamp_ns"]   = val->rgb_timestamp_ns;
        data["depth_timestamp_ns"] = val->depth_timestamp_ns;

        return data;
    }
};

// -----------------------------------------------------------------------
// VoiceQueryReader
// -----------------------------------------------------------------------

struct py_voice_query_reader {
    switchboard::reader<data_format::semantic_xr::voice_query>* reader_;
    uint64_t                                                    last_query_id_ = 0;

    pybind11::object get() {
        auto val = reader_->get_ro_nullable();
        if (!val)
            return pybind11::none();

        // Only return if this is a new query we haven't delivered yet
        if (val->query_id == last_query_id_)
            return pybind11::none();

        last_query_id_ = val->query_id;

        pybind11::dict d;
        d["query_id"]             = val->query_id;
        d["similarity_threshold"] = val->similarity_threshold;
        d["min_match_similarity"] = val->min_match_similarity;

        // PCM data — zero-copy with capsule
        pybind11::capsule base(new std::shared_ptr<const data_format::semantic_xr::voice_query>(val), [](void* p) {
            delete static_cast<std::shared_ptr<const data_format::semantic_xr::voice_query>*>(p);
        });
        d["pcm_data"] = pybind11::array_t<uint8_t>({static_cast<pybind11::ssize_t>(val->pcm_data.size())},
                                                   {static_cast<pybind11::ssize_t>(1)}, val->pcm_data.data(), base);

        return d;
    }
};

// -----------------------------------------------------------------------
// QueryResponseWriter
// -----------------------------------------------------------------------

struct py_query_response_writer {
    switchboard::writer<data_format::semantic_xr::query_response>* writer_;

    void put(uint64_t query_id, pybind11::list point_clouds, pybind11::list colors, float server_latency,
             const std::string& text_query) const {
        auto resp                     = std::make_shared<data_format::semantic_xr::query_response>();
        resp->query_id                = query_id;
        resp->server_query_processing = server_latency;
        resp->text_query              = text_query;
        resp->num_point_clouds        = static_cast<int32_t>(point_clouds.size());

        resp->point_clouds.reserve(point_clouds.size());
        for (auto& item : point_clouds) {
            auto                                  pc_dict = item.cast<pybind11::dict>();
            data_format::semantic_xr::point_cloud pc;
            pc.points     = pc_dict["points"].cast<std::vector<float>>();
            pc.num_points = static_cast<int32_t>(pc.points.size() / 3);
            pc.centroid   = pc_dict["centroid"].cast<std::vector<float>>();
            resp->point_clouds.push_back(std::move(pc));
        }

        resp->colors = colors.cast<std::vector<float>>();

        writer_->put(std::move(resp));
    }
};

inline void register_bindings(pybind11::module_& m) {
    pybind11::class_<py_semantic_data_reader>(m, "DnnInputReader").def("get", &py_semantic_data_reader::get);

    pybind11::class_<py_voice_query_reader>(m, "VoiceQueryReader")
        .def("get", &py_voice_query_reader::get, "Return the latest voice query as a dict, or None if no new query.");

    pybind11::class_<py_query_response_writer>(m, "QueryResponseWriter")
        .def("put", &py_query_response_writer::put, pybind11::arg("query_id"), pybind11::arg("point_clouds"),
             pybind11::arg("colors"), pybind11::arg("server_latency"), pybind11::arg("text_query"),
             "Write a query response to the switchboard.");
}
} // namespace ILLIXR