dataset_parser.h

Go to the documentation of this file.

/*
* ECVL - European Computer Vision Library
* Version: 1.0.0
* 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_PARSER_H_
#define ECVL_DATASET_PARSER_H_

#include "ecvl/core.h"
#include "ecvl/core/any.h"
#include "ecvl/core/filesystem.h"
#include "ecvl/core/optional.h"

#include <iostream>
#include <iterator>
#include <map>
#include <vector>
#include <regex>

#include "yaml-cpp/yaml.h"

// This allows to define strongly typed enums and convert them to int with just a + in front
#define UNSIGNED_ENUM_CLASS(name, ...) enum class name : unsigned { __VA_ARGS__ };\
inline constexpr unsigned operator+ (name const val) { return static_cast<unsigned>(val); }

namespace ecvl
{
UNSIGNED_ENUM_CLASS(SplitType, training, validation, test)


enum class Task
{
    classification,
    segmentation,
};

class Sample
{
public:
    std::vector<filesystem::path> location_;        
    optional<std::vector<int>> label_;              
    optional<filesystem::path> label_path_;         
    optional<std::map<int, std::string>> values_;   
    std::vector<int> size_;                         
    ecvl::Image LoadImage(ecvl::ColorType ctype = ecvl::ColorType::RGB, const bool& is_gt = false);
};

class Split
{
public:
    std::string split_name_;                
    optional<SplitType> split_type_;        
    std::vector<int> samples_indices_;      
    bool drop_last_ = false;                
    int num_batches_;                       
    int last_batch_;                        
    bool no_label_ = false;                 
    Split() = default;

    Split(const std::string& split_name, const std::vector<int>& samples_indices, const bool drop_last = false, const bool no_label = false)
        : split_name_{ split_name }, samples_indices_{ samples_indices }, drop_last_{ drop_last }, no_label_{ no_label }
    {
        if (split_name_ == "training") split_type_ = SplitType::training;
        else if (split_name_ == "validation") split_type_ = SplitType::validation;
        else if (split_name_ == "test") split_type_ = SplitType::test;
    }

    void SetNumBatches(int batch_size)
    {
        num_batches_ = drop_last_ ? vsize(samples_indices_) / batch_size : (vsize(samples_indices_) + batch_size - 1) / batch_size;
    }

    void SetLastBatch(int batch_size)
    {
        // last batch is the remainder of the number of samples of the split divided by the batch size.
        // if drop last is true or the remainder is 0, last batch is equal to the batch size.
        auto value = vsize(samples_indices_) % batch_size;
        last_batch_ = drop_last_ ? batch_size : (value == 0 ? batch_size : value);
    }
};

class Dataset
{
    std::map<std::string, int> features_map_;
    void DecodeImages(const YAML::Node& node, const filesystem::path& root_path, bool verify);
    void FindLabel(Sample& sample, const YAML::Node& n);
protected:
    std::vector<ecvl::Split>::iterator GetSplitIt(ecvl::any split);
    const int GetSplitIndex(ecvl::any split);
public:
    std::string name_ = "DeepHealth dataset";                               
    std::string description_ = "This is the DeepHealth example dataset!";   
    std::vector<std::string> classes_;                                      
    std::vector<std::string> features_;                                     
    std::vector<Sample> samples_;                                           
    std::vector<Split> split_;                                              
    int current_split_ = -1;                                                
    Task task_;                                                             
    Dataset() {}

    Dataset(const filesystem::path& filename, bool verify = false);

    /* Destructor */
    virtual ~Dataset() {}

    std::vector<int>& GetSplit(const ecvl::any& split = -1);

    void SetSplit(const ecvl::any& split);

    void Dump(const filesystem::path& file_path);

    std::vector<std::vector<filesystem::path>> GetLocations() const;

    // RegEx which matchs URLs
    static const std::regex url_regex_;
};
} // namespace ecvl

#endif // ECVL_DATASET_PARSER_H_