Style Fixed

mlpack · Jun 4, 2020 · e97e616 · e97e616
1 parent 93acea2
commit e97e616
Showing 1 changed file with 15 additions and 1 deletion.
diff --git a/dataloader/dataloader_impl.hpp b/dataloader/dataloader_impl.hpp
@@ -73,7 +73,7 @@ template<
 
       LoadObjectDetectionDataset(datasetMap[dataset].trainingAnnotationPath,
           datasetMap[dataset].trainingImagesPath, datasetMap[dataset].classes,
-          validRatio, shuffle,augmentations, augmentationProbability);
+          validRatio, shuffle, augmentations, augmentationProbability);
 
       // Load testing data if any. Most object detection dataset
       // have private evaluation servers.
@@ -84,6 +84,20 @@ template<
             datasetMap[dataset].imageHeight, datasetMap[dataset].imageDepth);
       }
     }
+    else if (datasetMap[dataset].datasetType == "image-classification")
+    {
+      LoadImageDatasetFromDirectory(datasetMap[dataset].trainingImagesPath,
+          datasetMap[dataset].imageWidth, datasetMap[dataset].imageHeight,
+          datasetMap[dataset].imageDepth, true, validRatio, shuffle,
+          augmentation, augmentationProbability);
+
+      if (datasetMap[dataset].testingImagesPath.length() > 0)
+      {
+        LoadAllImagesFromDirectory(datasetMap[dataset].testingImagesPath,
+            testFeatures, testLabels, datasetMap[dataset].imageWidth,
+            datasetMap[dataset].imageHeight, datasetMap[dataset].imageDepth);
+      }
+    }
 
     // Preprocess the dataset.
     datasetMap[dataset].PreProcess(trainFeatures, trainLabels,