lightly_edge_sdk.h

#ifndef LIGHTLY_EDGE_SDK_H
#define LIGHTLY_EDGE_SDK_H
 
#include "cxx.h"
#include "lightly_edge_cxx.rs.h"
#include "lightly_edge_error.h"
#include "lightly_edge_rs_bindings.h"
#include "lightly_edge_rs_bindings_ext.h"
#include "tl/expected.hpp"
#include <cassert>
#include <cstdint>
#include <cstdio>
#include <string>
#include <vector>
 
namespace lightly_edge_sdk {
 
const size_t BUFFER_MAX_LENGTH = 3000;
const uint32_t MAX_CLASSIFICATIONS_DEFAULT = 5;
const int CPP_SDK_ERROR_CODE = 30;
 
struct Frame {
  size_t width_;
 
  size_t height_;
 
  void *rgbImageData_;
 
  // NOLINTNEXTLINE(bugprone-easily-swappable-parameters)
  Frame(size_t width, size_t height, void *rgbImageData)
      : width_(width), height_(height), rgbImageData_(rgbImageData) {}
};
 
using lightly_edge_cxx::InferenceDeviceType;
using lightly_edge_cxx::LightlyEdgeConfig;
using lightly_edge_rs::AdaptiveDiversitySelectInfo;
using lightly_edge_rs::CStatus;
using lightly_edge_rs::DetectionSelectInfo;
using lightly_edge_rs::DiversitySelectInfo;
using lightly_edge_rs::PatchCoordsRel;
using lightly_edge_rs::SimilaritySelectInfo;
 
inline auto default_config() -> LightlyEdgeConfig {
  LightlyEdgeConfig config{};
  config.detector_config.object_detector_enable = false;
  config.detector_config.classifiers_enable = false;
  config.detector_config.max_classifications = MAX_CLASSIFICATIONS_DEFAULT;
  config.inference_device_type = InferenceDeviceType::Auto;
  return config;
}
 
// Converts CxxStatus to tl::unexpected.
inline auto cxx_status_to_unexpected(lightly_edge_cxx::CxxStatus &status)
    -> tl::unexpected<Error> {
  return tl::make_unexpected(
      lightly_edge_sdk::Error(status.code, std::string(status.message)));
}
 
struct SelectInfo {
  std::vector<DiversitySelectInfo> diversity;
 
  std::vector<SimilaritySelectInfo> similarity;
 
  std::vector<AdaptiveDiversitySelectInfo> adaptive_diversity;
 
  std::vector<DetectionSelectInfo> detection;
};
 
using lightly_edge_rs::ObjectDetection;
 
class LightlyEdge {
public:
  ~LightlyEdge() {
    if (this->lightly_edge != nullptr) {
      lightly_edge_rs::destroy(this->lightly_edge);
    }
  }
 
  LightlyEdge(LightlyEdge &&other) noexcept
      : lightly_edge(other.lightly_edge), lightly_edge_cxx(other.lightly_edge_cxx) {
    other.lightly_edge = nullptr;
    other.lightly_edge_cxx = nullptr;
  }
 
  auto operator=(LightlyEdge &&other) noexcept -> LightlyEdge & {
    if (this != &other) {
      if (this->lightly_edge != nullptr) {
        lightly_edge_rs::destroy(this->lightly_edge);
      }
      this->lightly_edge = other.lightly_edge;
      this->lightly_edge_cxx = other.lightly_edge_cxx;
      other.lightly_edge = nullptr;
      other.lightly_edge_cxx = nullptr;
    }
    return *this;
  }
 
  LightlyEdge(const LightlyEdge &other) = delete;
 
  auto operator=(const LightlyEdge &other) -> LightlyEdge & = delete;
 
  [[nodiscard]] static auto new_from_tar(const std::string &archive_path,
                                         LightlyEdgeConfig config) -> LightlyEdge {
    return value_or_throw(LightlyEdge::new_from_tar_noexcept(archive_path, config));
  }
 
  [[nodiscard]] static auto new_from_tar_noexcept(const std::string &archive_path,
                                                  LightlyEdgeConfig config) noexcept
      -> tl::expected<LightlyEdge, Error> {
    lightly_edge_cxx::LightlyEdge *lightly_edge = nullptr;
    lightly_edge_cxx::CxxStatus status =
        lightly_edge_cxx::new_from_tar(archive_path, config, &lightly_edge);
    if (status.code != lightly_edge_rs::SUCCESS_CODE) {
      return cxx_status_to_unexpected(status);
    }
    return LightlyEdge(lightly_edge);
  }
 
  [[nodiscard]] static auto new_from_directory(const std::string &model_directory,
                                               LightlyEdgeConfig config) {
    return value_or_throw(
        LightlyEdge::new_from_directory_noexcept(model_directory, config));
  }
 
  [[nodiscard]] static auto
  new_from_directory_noexcept(const std::string &model_directory,
                              LightlyEdgeConfig config) noexcept
      -> tl::expected<LightlyEdge, Error> {
    lightly_edge_cxx::LightlyEdge *lightly_edge = nullptr;
 
    auto status =
        lightly_edge_cxx::new_from_directory(model_directory, config, &lightly_edge);
    if (status.code != lightly_edge_rs::SUCCESS_CODE) {
      return cxx_status_to_unexpected(status);
    }
    return LightlyEdge(lightly_edge);
  }
 
  [[nodiscard]] auto embed_frame(const Frame &frame) const -> std::vector<float> {
    return value_or_throw(embed_frame_noexcept(frame));
  }
 
  [[nodiscard]] auto embed_frame_noexcept(const Frame &frame) const noexcept
      -> tl::expected<std::vector<float>, Error> {
    size_t embedding_dimension =
        lightly_edge_rs::embedding_dimension(this->lightly_edge);
 
    std::vector<float> embedding_buffer(embedding_dimension);
    auto status = lightly_edge_rs::embed_frame(
        this->lightly_edge, static_cast<uint8_t *>(frame.rgbImageData_),
        frame.width_ * frame.height_ * 3, frame.width_, frame.height_,
        embedding_buffer.data());
    if (status.code != lightly_edge_rs::SUCCESS_CODE) {
      return status_to_unexpected(&status);
    }
    // Moving the vector is more efficient than copying it.
    return std::move(embedding_buffer);
  }
 
  [[nodiscard]] auto embedding_dimension() const noexcept -> size_t {
    return lightly_edge_rs::embedding_dimension(this->lightly_edge);
  }
 
  [[nodiscard]] auto embed_texts(const std::vector<std::string> &texts) const
      -> std::vector<std::vector<float>> {
    return value_or_throw(embed_texts_noexcept(texts));
  }
 
  [[nodiscard]] auto
  embed_texts_noexcept(const std::vector<std::string> &texts) const noexcept
      -> tl::expected<std::vector<std::vector<float>>, Error> {
 
    rust::Vec<lightly_edge_cxx::Embedding> rust_embeddings{};
    // Call the Rust function to fill in the embeddings.
    auto status =
        lightly_edge_cxx::embed_texts(*this->lightly_edge_cxx, texts, rust_embeddings);
 
    // Check if the embedding was successful.
    if (status.code != lightly_edge_rs::SUCCESS_CODE) {
      return cxx_status_to_unexpected(status);
    }
 
    std::vector<std::vector<float>> embeddings;
    embeddings.resize(rust_embeddings.size());
    for (size_t i = 0; i < rust_embeddings.size(); ++i) {
      embeddings[i] = std::vector<float>(rust_embeddings[i].data.begin(),
                                         rust_embeddings[i].data.end());
    }
 
    // Return the populated embeddings vector.
    return embeddings;
  }
 
  void insert_into_embedding_database(const std::vector<float> &embedding) const {
    value_or_throw(insert_into_embedding_database_noexcept(embedding));
  }
 
  [[nodiscard]] auto insert_into_embedding_database_noexcept(
      const std::vector<float> &embedding) const noexcept -> tl::expected<void, Error> {
    auto status = lightly_edge_rs::insert_into_embedding_database(
        this->lightly_edge, embedding.data(), embedding.size());
    if (status.code != lightly_edge_rs::SUCCESS_CODE) {
      return status_to_unexpected(&status);
    }
    return {};
  }
 
  auto clear_embedding_database() const -> void {
    value_or_throw(clear_embedding_database_noexcept());
  }
 
  [[nodiscard]] auto clear_embedding_database_noexcept() const noexcept
      -> tl::expected<void, Error> {
    auto status = lightly_edge_cxx::clear_embedding_database(*this->lightly_edge_cxx);
    if (status.code != lightly_edge_rs::SUCCESS_CODE) {
      return cxx_status_to_unexpected(status);
    }
    return {};
  }
 
  [[nodiscard]] auto num_diversity_strategies() const noexcept -> size_t {
    return lightly_edge_rs::num_diversity_strategies(this->lightly_edge);
  }
 
  [[nodiscard]] auto num_similarity_strategies() const noexcept -> size_t {
    return lightly_edge_rs::num_similarity_strategies(this->lightly_edge);
  }
 
  [[nodiscard]] auto num_adaptive_diversity_strategies() const noexcept -> size_t {
    return lightly_edge_rs::num_adaptive_diversity_strategies(this->lightly_edge);
  }
 
  [[nodiscard]] auto num_detection_strategies() const noexcept -> size_t {
    return lightly_edge_rs::num_detection_strategies(this->lightly_edge);
  }
  auto register_diversity_strategy(float min_distance) const -> void {
    value_or_throw(register_diversity_strategy_noexcept(min_distance));
  }
 
  [[nodiscard]] auto
  register_diversity_strategy_noexcept(float min_distance) const noexcept
      -> tl::expected<void, Error> {
    auto status =
        lightly_edge_rs::register_diversity_strategy(this->lightly_edge, min_distance);
    if (status.code != lightly_edge_rs::SUCCESS_CODE) {
      return status_to_unexpected(&status);
    }
    return {};
  }
 
  auto register_similarity_strategy(const std::vector<float> &query_embedding,
                                    float max_distance) const -> void {
    value_or_throw(
        register_similarity_strategy_noexcept(query_embedding, max_distance));
  }
 
  [[nodiscard]] auto
  register_similarity_strategy_noexcept(const std::vector<float> &query_embedding,
                                        float max_distance) const noexcept
      -> tl::expected<void, Error> {
    auto status = lightly_edge_rs::register_similarity_strategy(
        this->lightly_edge, query_embedding.data(), query_embedding.size(),
        max_distance);
    if (status.code != lightly_edge_rs::SUCCESS_CODE) {
      return status_to_unexpected(&status);
    }
    return {};
  }
 
  auto register_adaptive_diversity_strategy(
      float target_ratio, size_t buffer_max_length = BUFFER_MAX_LENGTH) const -> void {
    value_or_throw(
        register_adaptive_diversity_strategy_noexcept(target_ratio, buffer_max_length));
  }
 
  [[nodiscard]] auto register_adaptive_diversity_strategy_noexcept(
      float target_ratio, size_t buffer_max_length = BUFFER_MAX_LENGTH) const noexcept
      -> tl::expected<void, Error> {
    auto status = lightly_edge_rs::register_adaptive_diversity_strategy(
        this->lightly_edge, target_ratio, buffer_max_length);
    if (status.code != lightly_edge_rs::SUCCESS_CODE) {
      return status_to_unexpected(&status);
    }
    return {};
  }
 
  auto register_detection_strategy(uint32_t class_id, int32_t subclass_id,
                                   float threshold) const -> void {
    value_or_throw(
        register_detection_strategy_noexcept(class_id, subclass_id, threshold));
  }
 
  [[nodiscard]] auto
  register_detection_strategy_noexcept(uint32_t class_id, int32_t subclass_id,
                                       float threshold) const noexcept
      -> tl::expected<void, Error> {
    auto status = lightly_edge_rs::register_detection_strategy(
        this->lightly_edge, class_id, subclass_id, threshold);
    if (status.code != lightly_edge_rs::SUCCESS_CODE) {
      return status_to_unexpected(&status);
    }
    return {};
  }
 
  void clear_strategies() { lightly_edge_rs::clear_strategies(this->lightly_edge); }
 
  [[nodiscard]] auto embed_patches(const Frame &frame) const
      -> std::vector<std::vector<float>> {
    return value_or_throw(embed_patches_noexcept(frame));
  }
 
  [[nodiscard]] auto embed_patches_noexcept(const Frame &frame) const noexcept
      -> tl::expected<std::vector<std::vector<float>>, Error> {
    size_t embedding_dimension =
        lightly_edge_rs::embedding_dimension(this->lightly_edge);
    size_t num_patches = lightly_edge_rs::num_patches(this->lightly_edge);
 
    // Prepare a vector of pointers to the data of each embedding vector.
    std::vector<std::vector<float>> embeddings(num_patches,
                                               std::vector<float>(embedding_dimension));
    std::vector<float *> embeddings_ptrs;
    embeddings_ptrs.reserve(num_patches);
    for (auto &embedding : embeddings) {
      embeddings_ptrs.push_back(embedding.data());
    }
 
    // Call the Rust function to fill in the embeddings.
    auto status = lightly_edge_rs::embed_patches(
        this->lightly_edge, static_cast<uint8_t *>(frame.rgbImageData_),
        frame.width_ * frame.height_ * 3, frame.width_, frame.height_,
        embeddings_ptrs.data());
 
    if (status.code != lightly_edge_rs::SUCCESS_CODE) {
      return status_to_unexpected(&status);
    }
 
    // Return the populated embeddings vector.
    return embeddings;
  }
 
  auto set_embedding_patches(const std::vector<PatchCoordsRel> &patches) const -> void {
    value_or_throw(set_embedding_patches_noexcept(patches));
  }
 
  [[nodiscard]] auto set_embedding_patches_noexcept(
      const std::vector<PatchCoordsRel> &patches) const noexcept
      -> tl::expected<void, Error> {
    auto status = lightly_edge_rs::set_embedding_patches(
        this->lightly_edge, patches.data(), patches.size());
    if (status.code != lightly_edge_rs::SUCCESS_CODE) {
      return status_to_unexpected(&status);
    }
    return {};
  }
 
  void reset_embedding_patches() {
    lightly_edge_cxx::reset_embedding_patches(*this->lightly_edge_cxx);
  }
 
  [[nodiscard]] auto should_select(const std::vector<float> &embedding,
                                   const std::vector<ObjectDetection> &detections) const
      -> SelectInfo {
    return value_or_throw(should_select_noexcept(embedding, detections));
  }
 
  [[nodiscard]] auto
  should_select_noexcept(const std::vector<float> &embedding,
                         const std::vector<ObjectDetection> &detections) const noexcept
      -> tl::expected<SelectInfo, Error> {
    // Initialize the output buffers.
    CStatus status = {};
    auto select_info_buffers = SelectInfoBuffers(this->lightly_edge);
    auto out_info = select_info_buffers.to_c_select_info();
 
    status = lightly_edge_rs::should_select(this->lightly_edge, embedding.data(),
                                            embedding.size(), detections.data(),
                                            detections.size(), &out_info);
    if (status.code != lightly_edge_rs::SUCCESS_CODE) {
      return status_to_unexpected(&status);
    }
 
    // Convert to the output C++ struct.
    auto select_info = select_info_buffers.to_select_info();
    return select_info;
  }
 
  [[nodiscard]] auto
  should_select_with_patches(const std::vector<std::vector<float>> &embeddings,
                             const std::vector<ObjectDetection> &detections) const
      -> SelectInfo {
    return value_or_throw(should_select_with_patches_noexcept(embeddings, detections));
  }
 
  [[nodiscard]] auto should_select_with_patches_noexcept(
      const std::vector<std::vector<float>> &embeddings,
      const std::vector<ObjectDetection> &detections) const noexcept
      -> tl::expected<SelectInfo, Error> {
    // Initialize the output buffers.
    CStatus status = {};
    auto select_info_buffers = SelectInfoBuffers(this->lightly_edge);
    auto out_info = select_info_buffers.to_c_select_info();
 
    // Get the embedding dimension. If there are no embeddings, the value does not
    // matter.
    // TODO(Michal, 12/2024): IMPORTANT! Calling this function with no embeddings will
    // fail in Rust. Decide how to handle this.
    size_t embedding_dim = embeddings.empty() ? 0 : embeddings[0].size();
 
    // Validate all embeddings have the same dimension.
    for (const auto &embedding : embeddings) {
      if (embedding.size() != embedding_dim) {
        return tl::make_unexpected(lightly_edge_sdk::Error(
            CPP_SDK_ERROR_CODE, "All patch embeddings must have the same dimension"));
      }
    }
 
    // Convert to a vector of pointers to the data of each embedding vector.
    std::vector<const float *> embeddings_ptrs;
    embeddings_ptrs.reserve(embeddings.size());
    for (const auto &embedding : embeddings) {
      embeddings_ptrs.push_back(embedding.data());
    }
 
    status = lightly_edge_rs::should_select_with_patches(
        this->lightly_edge, embeddings_ptrs.data(), embeddings_ptrs.size(),
        embedding_dim, detections.data(), detections.size(), &out_info);
    if (status.code != lightly_edge_rs::SUCCESS_CODE) {
      return status_to_unexpected(&status);
    }
 
    // Convert to the output C++ struct.
    auto select_info = select_info_buffers.to_select_info();
    return select_info;
  }
 
  // TODO(lukas 8/24): Add docstrings
  [[nodiscard]] auto detect_noexcept(const Frame &frame) const noexcept
      -> tl::expected<std::vector<ObjectDetection>, Error> {
    std::vector<ObjectDetection> detections;
    // TODO(lukas 7/25): 1000 is a constant for the current model, we should read this
    // from JSON
    const auto capacity = 1000;
    detections.resize(capacity);
    size_t detections_count = detections.size();
    auto status = lightly_edge_rs::detect(
        this->lightly_edge, static_cast<uint8_t *>(frame.rgbImageData_),
        frame.width_ * frame.height_ * 3, frame.width_, frame.height_,
        detections.data(), &detections_count);
    detections.resize(detections_count);
    if (status.code != lightly_edge_rs::SUCCESS_CODE) {
      return status_to_unexpected(&status);
    }
    return detections;
  }
 
  [[nodiscard]] auto detect(const Frame &frame) const -> std::vector<ObjectDetection> {
    return value_or_throw(detect_noexcept(frame));
  }
 
  [[nodiscard]] auto detection_class_labels() const noexcept
      -> std::vector<std::string> {
    // TODO(lukas 8/24): consider exposing the number of label classes from Rust (or
    // switch to cxx).
    const auto capacity = 1000;
    size_t classes_count = capacity;
    std::vector<const char *> cstr_class_labels(classes_count);
    auto status = lightly_edge_rs::detection_class_labels(
        this->lightly_edge, cstr_class_labels.data(), &classes_count);
    // class_labels() failing is our fault. We either provided an invalid pointer or the
    // array is too short. This is a bug, so we assert here.
    // Note: assert only runs when NDEBUG is undefined.
    assert(status.code == lightly_edge_rs::SUCCESS_CODE); // NOLINT
 
    // Convert C-strings to std::string's.
    std::vector<std::string> class_labels(classes_count);
    for (size_t i = 0; i < classes_count; ++i) {
      class_labels[i] = cstr_class_labels[i];
    }
    return class_labels;
  }
 
  [[nodiscard]] auto detection_subclass_labels(uint32_t class_id) const noexcept
      -> std::vector<std::string> {
    // TODO(lukas 8/24): consider exposing the number of label classes from Rust (or
    // switch to cxx).
    const auto capacity = 1000;
    size_t subclasses_count = capacity;
    std::vector<const char *> cstr_subclass_labels(subclasses_count);
    auto status = lightly_edge_rs::detection_subclass_labels(
        this->lightly_edge, class_id, cstr_subclass_labels.data(), &subclasses_count);
    // detection_subclass_labels() failing is our fault. We either provided an invalid
    // pointer or the array is too short. This is a bug, so we assert here. Note: assert
    // only runs when NDEBUG is undefined.
    assert(status.code == lightly_edge_rs::SUCCESS_CODE); // NOLINT
 
    // Convert C-strings to std::string's.
    std::vector<std::string> subclass_labels(subclasses_count);
    for (size_t i = 0; i < subclasses_count; ++i) {
      subclass_labels[i] = cstr_subclass_labels[i];
    }
    return subclass_labels;
  }
 
private:
  // Private constructor. Use `new_from_tar` to create a new instance.
  explicit LightlyEdge(lightly_edge_rs::LightlyEdge *lightly_edge_)
      : lightly_edge(lightly_edge_),
        lightly_edge_cxx(
            reinterpret_cast<lightly_edge_cxx::LightlyEdge *>(lightly_edge_)) {}
 
  explicit LightlyEdge(lightly_edge_cxx::LightlyEdge *lightly_edge_cxx_)
      : lightly_edge(
            reinterpret_cast<lightly_edge_rs::LightlyEdge *>(lightly_edge_cxx_)),
        lightly_edge_cxx(lightly_edge_cxx_) {}
  lightly_edge_cxx::LightlyEdge *lightly_edge_cxx{};
  lightly_edge_rs::LightlyEdge *lightly_edge{};
 
  class SelectInfoBuffers {
  public:
    explicit SelectInfoBuffers(lightly_edge_rs::LightlyEdge *lightly_edge)
        : out_diversity_size(lightly_edge_rs::num_diversity_strategies(lightly_edge)),
          out_similarity_size(lightly_edge_rs::num_similarity_strategies(lightly_edge)),
          out_adaptive_diversity_size(
              lightly_edge_rs::num_adaptive_diversity_strategies(lightly_edge)),
          out_detection_selection_size(
              lightly_edge_rs::num_detection_strategies(lightly_edge)) {
 
      resize_buffers_to_match_sizes();
    }
 
    [[nodiscard]] auto to_c_select_info()
        -> lightly_edge_rs::SelectionStrategiesSelectInfo {
      lightly_edge_rs::SelectionStrategiesSelectInfo out_info{};
      out_info.out_diversity_buffer = diversity.data();
      out_info.diversity_buffer_size = diversity.size();
      out_info.out_diversity_size = &out_diversity_size;
 
      out_info.out_similarity_buffer = similarity.data();
      out_info.similarity_buffer_size = similarity.size();
      out_info.out_similarity_size = &out_similarity_size;
 
      out_info.out_adaptive_diversity_buffer = adaptive_diversity.data();
      out_info.adaptive_diversity_buffer_size = adaptive_diversity.size();
      out_info.out_adaptive_diversity_size = &out_adaptive_diversity_size;
 
      out_info.out_detection_buffer = detection.data();
      out_info.detection_buffer_size = detection.size();
      out_info.out_detection_size = &out_detection_selection_size;
 
      return out_info;
    }
 
    [[nodiscard]] auto to_select_info() -> lightly_edge_sdk::SelectInfo {
      // Important: should_select* methods modify the out_*_size variables. We need to
      // resize the vectors to the correct size before moving the data to the SelectInfo
      // struct.
      resize_buffers_to_match_sizes();
 
      lightly_edge_sdk::SelectInfo select_info;
      select_info.diversity = std::move(diversity);
      select_info.similarity = std::move(similarity);
      select_info.adaptive_diversity = std::move(adaptive_diversity);
      select_info.detection = std::move(detection);
      return select_info;
    }
 
  private:
    size_t out_diversity_size;
    size_t out_similarity_size;
    size_t out_adaptive_diversity_size;
    size_t out_detection_selection_size;
    std::vector<lightly_edge_rs::DiversitySelectInfo> diversity;
    std::vector<lightly_edge_rs::SimilaritySelectInfo> similarity;
    std::vector<lightly_edge_rs::AdaptiveDiversitySelectInfo> adaptive_diversity;
    std::vector<lightly_edge_rs::DetectionSelectInfo> detection;
 
    void resize_buffers_to_match_sizes() {
      diversity.resize(out_diversity_size);
      similarity.resize(out_similarity_size);
      adaptive_diversity.resize(out_adaptive_diversity_size);
      detection.resize(out_detection_selection_size);
    }
  };
};
 
} // namespace lightly_edge_sdk
 
#endif // LIGHTLY_EDGE_SDK_H