support_eddl.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_SUPPORT_EDDL_H_
#define ECVL_SUPPORT_EDDL_H_

#include "ecvl/augmentations.h"
#include "ecvl/core/filesystem.h"
#include "ecvl/core/image.h"
#include "ecvl/dataset_parser.h"

#include <eddl/apis/eddl.h>

#include <algorithm>
#include <condition_variable>
#include <mutex>
#include <queue>
#include <thread>
#include <tuple>

namespace ecvl
{
#define ECVL_ERROR_AUG_DOES_NOT_EXIST throw std::runtime_error(ECVL_ERROR_MSG "Augmentation for this split does not exist");
#define ECVL_ERROR_WORKERS_LESS_THAN_ONE throw std::runtime_error(ECVL_ERROR_MSG "Dataset workers must be at least one");

void ImageToTensor(const Image& img, Tensor*& t);

void ImageToTensor(const Image& img, Tensor*& t, const int& offset);

void TensorToImage(const Tensor* t, Image& img);

void TensorToView(const Tensor* t, View<DataType::float32>& v);

class DatasetAugmentations
{
    std::vector<std::shared_ptr<Augmentation>> augs_;
public:
    DatasetAugmentations(const std::vector<std::shared_ptr<Augmentation>>& augs) : augs_(augs) {}

    // This makes a deep copy of the Augmentations
    DatasetAugmentations(const DatasetAugmentations& other)
    {
        for (const auto& a : other.augs_) {
            augs_.emplace_back(a ? a->Clone() : nullptr);
        }
    }

    // Getters: YAGNI

    bool Apply(const int split, Image& img, const Image& gt = Image())
    {
        // check if the augs for split st are provided
        try {
            if (augs_.at(split)) {
                augs_[split]->Apply(img, gt);
                return true;
            }
            return false;
        }
        catch (const std::out_of_range) {
            ECVL_ERROR_AUG_DOES_NOT_EXIST
        }
    }

    bool Apply(SplitType st, Image& img, const Image& gt = Image())
    {
        return Apply(+st, img, gt); // Magic + operator
    }

    bool IsEmpty() const
    {
        return augs_.empty();
    }
};

class ProducersConsumerQueue
{
    std::condition_variable cond_notempty_;     
    std::condition_variable cond_notfull_;      
    mutable std::mutex mutex_;                  
    std::queue<std::tuple<Sample, Tensor*, Tensor*>> cpq_;  
    unsigned max_size_;                         
    unsigned threshold_;                        
public:
    ProducersConsumerQueue() {}
    ProducersConsumerQueue(unsigned mxsz) : max_size_(mxsz), threshold_(max_size_ / 2) {}
    ProducersConsumerQueue(unsigned mxsz, unsigned thresh) : max_size_(mxsz), threshold_(thresh) {}

    void Push(const Sample& sample, Tensor* const image, Tensor* const label)
    {
        std::unique_lock<std::mutex> lock(mutex_);
        cond_notfull_.wait(lock, [this]() { return cpq_.size() < max_size_; });
        cpq_.push(std::make_tuple(sample, image, label));
        cond_notempty_.notify_one();
    }

    void Pop(Sample& sample, Tensor*& image, Tensor*& label)
    {
        std::unique_lock<std::mutex> lock(mutex_);
        cond_notempty_.wait(lock, [this]() { return !cpq_.empty(); });
        std::tie(sample, image, label) = cpq_.front();
        cpq_.pop();
        if (cpq_.size() < threshold_) {
            cond_notfull_.notify_all();
        }
    }

    bool IsFull() const
    {
        std::unique_lock<std::mutex> lock(mutex_);
        return cpq_.size() >= max_size_;
    }

    bool IsEmpty() const
    {
        std::unique_lock<std::mutex> lock(mutex_);
        return cpq_.empty();
    }

    size_t Length() const
    {
        std::unique_lock<std::mutex> lock(mutex_);
        return cpq_.size();
    }

    void SetSize(int max_size, int thresh = -1)
    {
        max_size_ = max_size;
        threshold_ = thresh != -1 ? thresh : max_size / 2;
    }

    void Clear()
    {
        std::unique_lock<std::mutex> lock(mutex_);

        // Remove residual samples and delete data
        while (!cpq_.empty()) {
            cpq_.pop();
        }
    }
};

class ThreadCounters
{
public:
    int counter_;   
    int min_, max_; 
    ThreadCounters(int min, int max) : counter_{ min }, min_{ min }, max_{ max } {}
    ThreadCounters(int counter, int min, int max) : counter_{ counter }, min_{ min }, max_{ max } {}
    void Reset() { counter_ = min_; }  
};

class DLDataset : public Dataset
{
    const unsigned processor_count_ = std::thread::hardware_concurrency();
protected:
    int batch_size_;                            
    std::vector<int> current_batch_;            
    ColorType ctype_;                           
    ColorType ctype_gt_;                        
    DatasetAugmentations augs_;                 
    unsigned num_workers_;                      
    ProducersConsumerQueue queue_;              
    std::pair< std::vector<int>, std::vector<int>> tensors_shape_; 
    std::vector<std::vector<ThreadCounters>> splits_tc_; 
    std::vector<std::thread> producers_;        
    bool active_ = false;                       
    std::mutex active_mutex_;                   
    static std::default_random_engine re_;      
    void InitTC(int split_index);

    void SetTensorsShape()
    {
        switch (task_) {
        case Task::classification:
            tensors_shape_ = std::make_pair<std::vector<int>, std::vector<int>>(
                { batch_size_, n_channels_, resize_dims_[0], resize_dims_[1] },
                { batch_size_, vsize(classes_) });
            break;
        case Task::segmentation:
            tensors_shape_ = std::make_pair<std::vector<int>, std::vector<int>>(
                { batch_size_, n_channels_, resize_dims_[0], resize_dims_[1] },
                { batch_size_, n_channels_gt_, resize_dims_[0], resize_dims_[1] });
            break;
        }
    }

public:
    int n_channels_;                            
    int n_channels_gt_ = -1;                    
    std::vector<int> resize_dims_;              
    DLDataset(const filesystem::path& filename,
        const int batch_size,
        const DatasetAugmentations& augs,
        const ColorType ctype = ColorType::RGB,
        const ColorType ctype_gt = ColorType::GRAY,
        const unsigned num_workers = 1,
        const double queue_ratio_size = 1.,
        const std::vector<bool>& drop_last = {},
        bool verify = false) :

        Dataset{ filename, verify },
        batch_size_{ batch_size },
        augs_(augs),
        num_workers_{ std::min(num_workers, processor_count_) },
        ctype_{ ctype },
        ctype_gt_{ ctype_gt },
        queue_{ static_cast<unsigned>(batch_size * queue_ratio_size * std::min(num_workers, processor_count_)) }
    {
        // resize current_batch_ to the number of splits and initialize it with 0
        current_batch_.resize(split_.size(), 0);

        // Initialize n_channels_
        Image tmp = samples_[0].LoadImage(ctype);
        n_channels_ = tmp.Channels();

        if (!split_.empty()) {
            current_split_ = 0;
            // Initialize resize_dims_ after that augmentations on the first image are performed
            if (augs_.IsEmpty()) {
                std::cout << ECVL_WARNING_MSG << "Augmentations are empty!" << std::endl;
            }
            else {
                while (!augs_.Apply(current_split_, tmp)) {
                    ++current_split_;
                }
            }
            auto y = tmp.channels_.find('y');
            auto x = tmp.channels_.find('x');
            assert(y != std::string::npos && x != std::string::npos);
            resize_dims_.insert(resize_dims_.begin(), { tmp.dims_[y],tmp.dims_[x] });

            // Initialize n_channels_gt_ if exists
            if (samples_[0].label_path_ != nullopt) {
                n_channels_gt_ = samples_[0].LoadImage(ctype_gt_, true).Channels();
            }
        }
        else {
            std::cout << ECVL_WARNING_MSG << "Missing splits in the dataset file." << std::endl;
        }

        // Set drop_last parameter for each split
        if (!drop_last.empty()) {
            if (vsize(drop_last) == vsize(split_)) {
                for (int i = 0; i < vsize(drop_last); ++i) {
                    split_[i].drop_last_ = drop_last[i];
                }
            }
            else {
                std::cout << ECVL_WARNING_MSG << "drop_last is not empty but the provided size is different from the size of the splits. The default value 'false' is set for all the splits" << std::endl;
            }
        }

        // Initialize num_batches, last_batch and the ThreadCounters for each split
        auto s_index = 0;
        splits_tc_ = std::vector<std::vector<ThreadCounters>>(vsize(split_));
        for (auto& s : split_) {
            s.SetNumBatches(batch_size_);
            s.SetLastBatch(batch_size_);

            InitTC(s_index);
            ++s_index;
        }

        SetTensorsShape();
    }

    void ResetBatch(const ecvl::any& split = -1, bool shuffle = false);

    void ResetAllBatches(bool shuffle = false);

    void LoadBatch(Tensor*& images, Tensor*& labels);

    void LoadBatch(Tensor*& images);

    static void SetSplitSeed(unsigned seed) { re_.seed(seed); }

    void SetBatchSize(int bs);

    virtual void ProduceImageLabel(DatasetAugmentations& augs, Sample& elem);

    void ThreadFunc(int thread_index);

    std::tuple<std::vector<Sample>, std::shared_ptr<Tensor>, std::shared_ptr<Tensor>> GetBatch();

    void Start(int split_index = -1);

    void Stop();

    auto GetQueueSize() const { return queue_.Length(); };

    void SetAugmentations(const DatasetAugmentations& da);

    const int GetNumBatches(const ecvl::any& split = -1);

    void ToTensorPlane(const std::vector<int>& label, Tensor*& tensor);

    void SetWorkers(const unsigned num_workers)
    {
        if (num_workers < 0) {
            ECVL_ERROR_WORKERS_LESS_THAN_ONE
        }

        num_workers_ = std::min(num_workers, processor_count_);
        for (int i = 0; i < vsize(split_); ++i) {
            InitTC(i);
        }
    }

    void SetNumChannels(const int n_channels, const int n_channels_gt = 1)
    {
        n_channels_ = n_channels;
        n_channels_gt_ = n_channels_gt;
        SetTensorsShape();
    }
};

Image MakeGrid(Tensor*& t, int cols = 8, bool normalize = false);

} // namespace ecvl

#endif // ECVL_SUPPORT_EDDL_H_