fix usedImageNum init bug in quantization tool

alibaba · Aug 23, 2019 · 27f45da · 27f45da
1 parent a92cefd
commit 27f45da
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 5 deletions.
diff --git a/tools/quantization/Helper.cpp b/tools/quantization/Helper.cpp
@@ -26,7 +26,7 @@ bool Helper::fileExist(const std::string& file) {
     return stat(file.c_str(), &buffer) == 0;
 }
 
-void Helper::readImages(std::vector<std::string>& images, const std::string& filePath, const int usedImageNum) {
+void Helper::readImages(std::vector<std::string>& images, const std::string& filePath, int* usedImageNum) {
     DIR* root = opendir(filePath.c_str());
     int count = 0;
     if (root == NULL) {
@@ -40,11 +40,11 @@ void Helper::readImages(std::vector<std::string>& images, const std::string& fil
             if (fileExist(fileName)) {
                 // std::cout << "==> " << fileName << std::endl;
                 // DLOG(INFO) << fileName;
-                if (usedImageNum <= 0) {
+                if (*usedImageNum == 0) {
                     // use all images in the folder
                     images.push_back(fileName);
                     count++;
-                } else if (count < usedImageNum) {
+                } else if (count < *usedImageNum) {
                     // use usedImageNum images
                     images.push_back(fileName);
                     count++;
@@ -55,6 +55,9 @@ void Helper::readImages(std::vector<std::string>& images, const std::string& fil
         }
         ent = readdir(root);
     }
+    if (*usedImageNum == 0) {
+        *usedImageNum = count;
+    }
     DLOG(INFO) << "used image num: " << images.size();
 }
 

diff --git a/tools/quantization/Helper.hpp b/tools/quantization/Helper.hpp
@@ -26,7 +26,7 @@ class Helper {
     static std::set<std::string> weightQuantizeMethod;
 
     static bool fileExist(const std::string& file);
-    static void readImages(std::vector<std::string>& images, const std::string& filePath, const int usedImageNum);
+    static void readImages(std::vector<std::string>& images, const std::string& filePath, int *usedImageNum);
     static void preprocessInput(MNN::CV::ImageProcess* pretreat, int targetWidth, int targetHeight,
                                 const std::string& inputImageFileName, MNN::Tensor* input);
     static void invertData(float* dst, const float* src, int size);

diff --git a/tools/quantization/calibration.cpp b/tools/quantization/calibration.cpp
@@ -60,6 +60,7 @@ Calibration::Calibration(MNN::NetT* model, uint8_t* modelBuffer, const int buffe
 
     config.sourceFormat = RGBA;
     std::string imagePath;
+    _imageNum = 0;
     {
         if (picObj.HasMember("mean")) {
             auto mean = picObj["mean"].GetArray();
@@ -112,7 +113,7 @@ Calibration::Calibration(MNN::NetT* model, uint8_t* modelBuffer, const int buffe
     _process = process;
 
     // read images file names
-    Helper::readImages(_imgaes, imagePath.c_str(), _imageNum);
+    Helper::readImages(_imgaes, imagePath.c_str(), &_imageNum);
 
     _initMNNSession(modelBuffer, bufferSize, channles);
     _initMaps();