diff --git a/augmentation/augmentation.hpp b/augmentation/augmentation.hpp index 4f538747..247013c6 100644 --- a/augmentation/augmentation.hpp +++ b/augmentation/augmentation.hpp @@ -36,14 +36,14 @@ class Augmentation /** * Constructor for augmentation class. * - * @param augmentation List of strings containing one of the supported + * @param augmentations List of strings containing one of the supported * augmentation. * @param augmentationProbability Probability of applying augmentation on * the dataset. * NOTE : This doesn't apply to augmentations * such as resize. */ - Augmentation(const std::vector& augmentation, + Augmentation(const std::vector& augmentations, const double augmentationProbability); /** diff --git a/augmentation/augmentation_impl.hpp b/augmentation/augmentation_impl.hpp index 75534422..0efbfb87 100644 --- a/augmentation/augmentation_impl.hpp +++ b/augmentation/augmentation_impl.hpp @@ -21,4 +21,19 @@ Augmentation::Augmentation() : { // Nothing to do here. } + +Augmentation::Augmentation(const std::vector& augmentations, + const double augmentationProbability) : + augmentations(augmentations), + augmentationProbability(augmentationProbability) +{ + // Sort the vector to place resize parameter to the front of the string. + // This prevents constant look ups for resize. + sort(this->augmentations.begin(), this->augmentations.end(), []( + std::string& str1, std::string& str2) + { + return str1.find("resize") != std::string::npos; + }); +} + #endif \ No newline at end of file