Yolo.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2025 Povl Filip Sonne-Frederiksen
//
// SPDX-License-Identifier: GPL-3.0-or-later
 
#pragma once
#include "reusex/vision/IModel.hpp"
#include "reusex/vision/libtorch/Data.hpp"
 
#include <torch/script.h>
#include <torch/torch.h>
 
#include <filesystem>
#include <memory>
#include <span>
#include <vector>
 
namespace reusex::vision::libtorch {

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

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

  std::vector<IDataset::Pair>
  forward(const std::span<IDataset::Pair> &input) override;
 
    private:
  torch::Tensor
  images_to_tensor(const std::vector<const LibTorchData *> &inputs) const;
 
  torch::jit::script::Module model_;
  torch::Device device_;
};
 
} // namespace reusex::vision::libtorch