image_folder_dataset.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <string>
#include <vector>
 
#include "pulsatrix/dataset.hpp"
#include "pulsatrix/device_backend.hpp"
 
namespace pulsatrix {
 
class ImageFolderDataset : public Dataset {
public:
    ImageFolderDataset(const std::string& root_dir, DeviceBackend* backend);
 
    [[nodiscard]] int64_t size() const override;
 
    [[nodiscard]] Sample get(int64_t index) const override;
 
    [[nodiscard]] const std::vector<std::string>& classes() const { return classes_; }
 
private:
    struct Entry {
        std::string path;
        int64_t label;
    };
 
    std::vector<Entry> entries_;
    std::vector<std::string> classes_;
    DeviceBackend* backend_;
};
 
}  // namespace pulsatrix