Fix bug with batch size = 1

Author: Nuzhny007

example/main.cpp

@@ -37,10 +37,11 @@ const char* keys =
 
 int main(int argc, char** argv)
 {
-    Help();
-
     cv::CommandLineParser parser(argc, argv, keys);
 
+    Help();
+    parser.printMessage();
+
     bool useOCL = parser.get<int>("gpu") != 0;
     cv::ocl::setUseOpenCL(useOCL);
     std::cout << (cv::ocl::useOpenCL() ? "OpenCL is enabled" : "OpenCL not used") << std::endl;

src/Detector/YoloDarknetDetector.cpp

@@ -101,6 +101,7 @@ void YoloDarknetDetector::Detect(const cv::UMat& colorFrame)
 	else
 	{
         std::vector<cv::Rect> crops = GetCrops(m_maxCropRatio, m_netSize, colorMat.size());
+		std::cout << "Image on " << crops.size() << " crops with size " << crops.front().size() << ", input size " << m_netSize << ", batch " << m_batchSize << ", frame " << colorMat.size() << std::endl;
         regions_t tmpRegions;
 		if (m_batchSize > 1)
 		{
@@ -306,7 +307,9 @@ void YoloDarknetDetector::Detect(const std::vector<cv::UMat>& frames, std::vecto
 {
 	if (frames.size() == 1)
 	{
-		Detect(frames[0].getMat(cv::ACCESS_READ), regions[0]);
+		Detect(frames[0]);
+		regions[0] = m_regions;
+
 	}
 	else
 	{

src/Detector/YoloTensorRTDetector.cpp

@@ -143,6 +143,7 @@ void YoloTensorRTDetector::Detect(const cv::UMat& colorFrame)
     else
     {
         std::vector<cv::Rect> crops = GetCrops(m_maxCropRatio, m_detector->get_input_size(), colorMat.size());
+		std::cout << "Image on " << crops.size() << " crops with size " << crops.front().size() << ", input size " << m_detector->get_input_size() << ", batch " << m_batchSize << ", frame " << colorMat.size() << std::endl;
         regions_t tmpRegions;
 		std::vector<cv::Mat> batch;
 		batch.reserve(m_batchSize);
@@ -189,22 +190,30 @@ void YoloTensorRTDetector::Detect(const cv::UMat& colorFrame)
 ///
 void YoloTensorRTDetector::Detect(const std::vector<cv::UMat>& frames, std::vector<regions_t>& regions)
 {
-	std::vector<cv::Mat> batch;
-	for (const auto& frame : frames)
-	{
-		batch.emplace_back(frame.getMat(cv::ACCESS_READ));
-	}
+    if (frames.size() == 1)
+    {
+        Detect(frames.front());
+        regions[0].assign(std::begin(m_regions), std::end(m_regions));
+    }
+    else
+    {
+        std::vector<cv::Mat> batch;
+        for (const auto& frame : frames)
+        {
+            batch.emplace_back(frame.getMat(cv::ACCESS_READ));
+        }
 
-	std::vector<tensor_rt::BatchResult> detects;
-	m_detector->detect(batch, detects);
-	for (size_t i = 0; i < detects.size(); ++i)
-	{
-		const tensor_rt::BatchResult& dets = detects[i];
-		for (const tensor_rt::Result& bbox : dets)
-		{
-			if (m_classesWhiteList.empty() || m_classesWhiteList.find(T2T(bbox.id)) != std::end(m_classesWhiteList))
-				regions[i].emplace_back(bbox.rect, T2T(bbox.id), bbox.prob);
-		}
-	}
-	m_regions.assign(std::begin(regions.back()), std::end(regions.back()));
+        std::vector<tensor_rt::BatchResult> detects;
+        m_detector->detect(batch, detects);
+        for (size_t i = 0; i < detects.size(); ++i)
+        {
+            const tensor_rt::BatchResult& dets = detects[i];
+            for (const tensor_rt::Result& bbox : dets)
+            {
+                if (m_classesWhiteList.empty() || m_classesWhiteList.find(T2T(bbox.id)) != std::end(m_classesWhiteList))
+                    regions[i].emplace_back(bbox.rect, T2T(bbox.id), bbox.prob);
+            }
+        }
+        m_regions.assign(std::begin(regions.back()), std::end(regions.back()));
+    }
 }