Sam3.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2026 Povl Filip Sonne-Frederiksen
//
// SPDX-License-Identifier: GPL-3.0-or-later
 
#pragma once
#include "reusex/vision/IModel.hpp"
#include "reusex/vision/common/object.hpp"
#include "reusex/vision/onnx/Sam3Data.hpp"
 
#include <onnxruntime_cxx_api.h>
#include <tokenizers_cpp.h>
 
#include <filesystem>
#include <memory>
#include <span>
#include <string>
#include <unordered_map>
#include <vector>
 
namespace reusex::vision::onnx {

class ONNXSam3 : public IModel {
    public:
  explicit ONNXSam3(const std::filesystem::path &model_dir,
                    bool use_cuda = false);

  static std::unique_ptr<ONNXSam3>
  create(const std::filesystem::path &model_dir, bool use_cuda = false);

  std::vector<IDataset::Pair>
  forward(const std::span<IDataset::Pair> &input) override;
 
    private:
  using InferResult = common::object::DetectionBoxArray;

  std::pair<std::vector<float>, cv::Mat> preprocess(const cv::Mat &image) const;

  std::pair<std::array<int64_t, 32>, std::array<int64_t, 32>>
  tokenize(const std::string &text);

  cv::Mat infer_single(const ONNXSam3Data &sam3_data);
 
  // ONNX Runtime state
  Ort::Env env_;
  Ort::SessionOptions session_options_;
  std::unique_ptr<Ort::Session> vision_encoder_;
  std::unique_ptr<Ort::Session> text_encoder_;
  std::unique_ptr<Ort::Session> decoder_;
 
  // Tokenizer
  std::unique_ptr<tokenizers::Tokenizer> tokenizer_;
 
  // Model parameters
  int input_size_ = 1008;
  int num_queries_ = 200;
  int mask_height_ = 288;
  int mask_width_ = 288;
 
  // Tokenizer cache: text -> (input_ids, attention_mask)
  std::unordered_map<
      std::string, std::pair<std::array<int64_t, 32>, std::array<int64_t, 32>>>
      text_cache_;
 
  // Memory allocator
  Ort::MemoryInfo memory_info_;
};
 
} // namespace reusex::vision::onnx