New features and fixes

nvdsinfer_custom_impl_Yolo/calibrator.cpp

@@ -8,17 +8,18 @@
 #include <fstream>
 #include <iterator>
 
-Int8EntropyCalibrator2::Int8EntropyCalibrator2(const int& batchsize, const int& channels, const int& height,
-    const int& width, const int& letterbox, const std::string& imgPath,
-    const std::string& calibTablePath) : batchSize(batchsize), inputC(channels), inputH(height), inputW(width),
-    letterBox(letterbox), calibTablePath(calibTablePath), imageIndex(0)
+Int8EntropyCalibrator2::Int8EntropyCalibrator2(const int& batchSize, const int& channels, const int& height, const int& width,
+    const float& scaleFactor, const float* offsets, const std::string& imgPath, const std::string& calibTablePath) :
+    batchSize(batchSize), inputC(channels), inputH(height), inputW(width), scaleFactor(scaleFactor), offsets(offsets),
+    calibTablePath(calibTablePath), imageIndex(0)
 {
-  inputCount = batchsize * channels * height * width;
+  inputCount = batchSize * channels * height * width;
   std::fstream f(imgPath);
   if (f.is_open()) {
       std::string temp;
-      while (std::getline(f, temp))
+      while (std::getline(f, temp)) {
         imgPaths.push_back(temp);
+      }
   }
   batchData = new float[inputCount];
   CUDA_CHECK(cudaMalloc(&deviceInput, inputCount * sizeof(float)));
@@ -27,8 +28,9 @@ Int8EntropyCalibrator2::Int8EntropyCalibrator2(const int& batchsize, const int&
 Int8EntropyCalibrator2::~Int8EntropyCalibrator2()
 {
   CUDA_CHECK(cudaFree(deviceInput));
-  if (batchData)
+  if (batchData) {
     delete[] batchData;
+  }
 }
 
 int
@@ -40,24 +42,33 @@ Int8EntropyCalibrator2::getBatchSize() const noexcept
 bool
 Int8EntropyCalibrator2::getBatch(void** bindings, const char** names, int nbBindings) noexcept
 {
-  if (imageIndex + batchSize > uint(imgPaths.size()))
+  if (imageIndex + batchSize > uint(imgPaths.size())) {
     return false;
+  }
 
   float* ptr = batchData;
   for (size_t i = imageIndex; i < imageIndex + batchSize; ++i) {
-    cv::Mat img = cv::imread(imgPaths[i], cv::IMREAD_COLOR);
-    std::vector<float> inputData = prepareImage(img, inputC, inputH, inputW, letterBox);
+    cv::Mat img = cv::imread(imgPaths[i]);
+    if (img.empty()){
+      std::cerr << "Failed to read image for calibration" << std::endl;
+      return false;
+    }
+  
+    std::vector<float> inputData = prepareImage(img, inputC, inputH, inputW, scaleFactor, offsets);
 
-    int len = (int) (inputData.size());
+    size_t len = inputData.size();
     memcpy(ptr, inputData.data(), len * sizeof(float));
-
     ptr += inputData.size();
+
     std::cout << "Load image: " << imgPaths[i] << std::endl;
-    std::cout << "Progress: " << (i + 1)*100. / imgPaths.size() << "%" << std::endl;
+    std::cout << "Progress: " << (i + 1) * 100. / imgPaths.size() << "%" << std::endl;
   }
+
   imageIndex += batchSize;
+
   CUDA_CHECK(cudaMemcpy(deviceInput, batchData, inputCount * sizeof(float), cudaMemcpyHostToDevice));
   bindings[0] = deviceInput;
+
   return true;
 }
 
@@ -67,8 +78,9 @@ Int8EntropyCalibrator2::readCalibrationCache(std::size_t &length) noexcept
   calibrationCache.clear();
   std::ifstream input(calibTablePath, std::ios::binary);
   input >> std::noskipws;
-  if (readCache && input.good())
+  if (readCache && input.good()) {
     std::copy(std::istream_iterator<char>(input), std::istream_iterator<char>(), std::back_inserter(calibrationCache));
+  }
   length = calibrationCache.size();
   return length ? calibrationCache.data() : nullptr;
 }
@@ -81,43 +93,24 @@ Int8EntropyCalibrator2::writeCalibrationCache(const void* cache, std::size_t len
 }
 
 std::vector<float>
-prepareImage(cv::Mat& img, int input_c, int input_h, int input_w, int letter_box)
+prepareImage(cv::Mat& img, int input_c, int input_h, int input_w, float scaleFactor, const float* offsets)
 {
   cv::Mat out;
+
+  cv::cvtColor(img, out, cv::COLOR_BGR2RGB);
+
   int image_w = img.cols;
   int image_h = img.rows;
-  if (image_w != input_w || image_h != input_h) {
-    if (letter_box == 1) {
-      float ratio_w = (float) image_w / (float) input_w;
-      float ratio_h = (float) image_h / (float) input_h;
-      if (ratio_w > ratio_h) {
-        int new_width = input_w * ratio_h;
-        int x = (image_w - new_width) / 2;
-        cv::Rect roi(abs(x), 0, new_width, image_h);
-        out = img(roi);
-      }
-      else if (ratio_w < ratio_h) {
-        int new_height = input_h * ratio_w;
-        int y = (image_h - new_height) / 2;
-        cv::Rect roi(0, abs(y), image_w, new_height);
-        out = img(roi);
-      }
-      else
-        out = img;
-      cv::resize(out, out, cv::Size(input_w, input_h), 0, 0, cv::INTER_CUBIC);
-    }
-    else {
-      cv::resize(img, out, cv::Size(input_w, input_h), 0, 0, cv::INTER_CUBIC);
-    }
-    cv::cvtColor(out, out, cv::COLOR_BGR2RGB);
-  }
-  else
-    cv::cvtColor(img, out, cv::COLOR_BGR2RGB);
 
-  if (input_c == 3)
-    out.convertTo(out, CV_32FC3, 1.0 / 255.0);
-  else
-    out.convertTo(out, CV_32FC1, 1.0 / 255.0);
+  if (image_w != input_w || image_h != input_h) {
+    float resizeFactor = std::max(input_w / (float) image_w, input_h / (float) img.rows);
+    cv::resize(out, out, cv::Size(0, 0), resizeFactor, resizeFactor, cv::INTER_CUBIC);
+    cv::Rect crop(cv::Point(0.5 * (out.cols - input_w), 0.5 * (out.rows - input_h)), cv::Size(input_w, input_h));
+    out = out(crop);
+  }
+
+  out.convertTo(out, CV_32F, scaleFactor);
+  cv::subtract(out, cv::Scalar(offsets[2] / 255, offsets[1] / 255, offsets[0] / 255), out, cv::noArray(), -1);
 
   std::vector<cv::Mat> input_channels(input_c);
   cv::split(out, input_channels);