dataset_generator.h

Go to the documentation of this file.

/*
* ECVL - European Computer Vision Library
* Version: 0.3.4
* copyright (c) 2021, Università degli Studi di Modena e Reggio Emilia (UNIMORE), AImageLab
* Authors:
*    Costantino Grana (costantino.grana@unimore.it)
*    Federico Bolelli (federico.bolelli@unimore.it)
*    Michele Cancilla (michele.cancilla@unimore.it)
*    Laura Canalini (laura.canalini@unimore.it)
*    Stefano Allegretti (stefano.allegretti@unimore.it)
* All rights reserved.
*/

#ifndef ECVL_DATASET_GENERATOR_H_
#define ECVL_DATASET_GENERATOR_H_

#include <ecvl/core/filesystem.h>
#include <ecvl/dataset_parser.h>

namespace ecvl
{
class GenerateDataset
{
public:
    const filesystem::path dataset_root_directory_;  
    std::vector<std::string> splits_;        
    std::vector<int> num_samples_;           
    ecvl::Dataset d_;                        
    GenerateDataset(const filesystem::path& dataset_root_directory) :
        dataset_root_directory_(dataset_root_directory)
    {
        for (auto& p : filesystem::directory_iterator(dataset_root_directory_)) {
            std::string tmp = p.path().stem().string();

            // Check if split folders exist
            if (tmp == "training" || tmp == "validation" || tmp == "test") {
                splits_.emplace_back(tmp);
            }
        }
        num_samples_.resize(splits_.size());

        d_.name_ = dataset_root_directory_.stem().string();
        d_.description_ = dataset_root_directory_.stem().string();
    }

    virtual ~GenerateDataset() = default;

    void LoadImagesAndSplits();

    Dataset GetDataset() { return d_; }

    virtual int LoadSplitImages(const filesystem::path& split) = 0;
};

class GenerateSegmentationDataset : public GenerateDataset
{
public:
    filesystem::path suffix_;  
    filesystem::path gt_name_; 
    GenerateSegmentationDataset(const filesystem::path& dataset_root_directory, filesystem::path suffix = "", filesystem::path gt_name = "") :
        GenerateDataset{ dataset_root_directory },
        suffix_{ suffix },
        gt_name_{ gt_name }
    {
        LoadImagesAndSplits();
    }

    virtual int LoadSplitImages(const filesystem::path& split) override;
};

class GenerateClassificationDataset : public GenerateDataset
{
public:

    GenerateClassificationDataset(const filesystem::path& dataset_root_directory) : GenerateDataset{ dataset_root_directory }
    {
        filesystem::path tmp;

        if (splits_.empty()) {
            tmp = dataset_root_directory_;
        }
        else {
            // use the first split available to list all the classes
            tmp = dataset_root_directory_ / splits_[0];
        }

        for (auto& p : filesystem::directory_iterator(tmp)) {
            if (filesystem::is_directory(p.path())) {
                d_.classes_.push_back(p.path().stem().string());
            }
        }

        LoadImagesAndSplits();
    }

    virtual int LoadSplitImages(const filesystem::path& split) override;
};
} // namespace ecvl

#endif // ECVL_DATASET_GENERATOR_H_