Dataset.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 <opencv2/core.hpp>
#include <torch/torch.h>
 
#include <filesystem>
#include <memory>
#include <vector>
 
// Forward declaration
namespace ReUseX::io {
class RTABMapDatabase;
}
 
namespace ReUseX::vision::libtorch {
class TorchDataset : public torch::data::datasets::Dataset<TorchDataset> {
  using Example = torch::data::Example<>;
 
    public:
  TorchDataset(std::filesystem::path dbPath = "");

  Example get(size_t index);

  torch::optional<size_t> size() const;

  void save(std::vector<cv::Mat> imgs, torch::Tensor index);
 
    private:
  std::shared_ptr<io::RTABMapDatabase> db_;
  std::vector<int> ids_;
};
} // namespace ReUseX::vision::libtorch