From 9618d5fc7fc158eea87f8300d35b1c80b1f388e2 Mon Sep 17 00:00:00 2001
From: Micaela Verucchi <micaelaverucchi@gmail.com>
Date: Tue, 19 May 2020 16:33:45 +0200
Subject: [PATCH] Fix original size

Signed-off-by: Micaela Verucchi <micaelaverucchi@gmail.com>
---
 include/tkDNN/DetectionNN.h | 5 +++--
 src/CenternetDetection.cpp  | 2 +-
 src/MobilenetDetection.cpp  | 4 ++--
 src/Yolo3Detection.cpp      | 4 ++--
 4 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/include/tkDNN/DetectionNN.h b/include/tkDNN/DetectionNN.h
index fc9a137..2e64be8 100644
--- a/include/tkDNN/DetectionNN.h
+++ b/include/tkDNN/DetectionNN.h
@@ -30,7 +30,7 @@ class DetectionNN {
         tk::dnn::NetworkRT *netRT = nullptr;
         dnnType *input_d;
 
-        cv::Size originalSize;
+        std::vector<cv::Size> originalSize;
 
         cv::Scalar colors[256];
 
@@ -103,13 +103,14 @@ class DetectionNN {
             if(cur_batches > nBatches)
                 FatalError("A batch size greater than nBatches cannot be used");
 
+            originalSize.clear();
             if(VERBOSE) printCenteredTitle(" TENSORRT detection ", '=', 30); 
             {
                 TIMER_START
                 for(int bi=0; bi<cur_batches;++bi){
                     if(!frames[bi].data)
                         FatalError("No image data feed to detection");
-                    originalSize = frames[bi].size();
+                    originalSize.push_back(frames[bi].size());
                     preprocess(frames[bi], bi);    
                 }
                 TIMER_STOP
diff --git a/src/CenternetDetection.cpp b/src/CenternetDetection.cpp
index 5d01cc3..9d8df38 100644
--- a/src/CenternetDetection.cpp
+++ b/src/CenternetDetection.cpp
@@ -127,7 +127,7 @@ void CenternetDetection::preprocess(cv::Mat &frame, const int bi){
     // auto start_t = std::chrono::steady_clock::now();
     // auto step_t = std::chrono::steady_clock::now();
     // auto end_t = std::chrono::steady_clock::now();
-    cv::Size sz = originalSize;
+    cv::Size sz = originalSize[bi];
     // std::cout<<"image: "<<sz.width<<", "<<sz.height<<std::endl;
     cv::Size sz_old;
     float scale = 1.0;
diff --git a/src/MobilenetDetection.cpp b/src/MobilenetDetection.cpp
index 4289d96..c905fea 100644
--- a/src/MobilenetDetection.cpp
+++ b/src/MobilenetDetection.cpp
@@ -256,8 +256,8 @@ void MobilenetDetection::postprocess(const int bi, const bool mAP){
     checkCuda(cudaMemcpy(locations_h, rt_out[1], N_COORDS * nPriors * sizeof(float), cudaMemcpyDeviceToHost));
     convert_locatios_to_boxes_and_center();
 
-    int width =  originalSize.width;
-    int height =  originalSize.height;
+    int width =  originalSize[bi].width;
+    int height =  originalSize[bi].height;
 
     float *conf_per_class;
     for (int i = 1; i < classes; i++){
diff --git a/src/Yolo3Detection.cpp b/src/Yolo3Detection.cpp
index fc5e1d0..c76af20 100644
--- a/src/Yolo3Detection.cpp
+++ b/src/Yolo3Detection.cpp
@@ -94,8 +94,8 @@ void Yolo3Detection::postprocess(const int bi, const bool mAP){
     for(int i=0; i<netRT->pluginFactory->n_yolos; i++) 
         rt_out[i] = (dnnType*)netRT->buffersRT[i+1] + netRT->buffersDIM[i+1].tot()*bi;
 
-    float x_ratio =  float(originalSize.width) / float(netRT->input_dim.w);
-    float y_ratio =  float(originalSize.height) / float(netRT->input_dim.h);
+    float x_ratio =  float(originalSize[bi].width) / float(netRT->input_dim.w);
+    float y_ratio =  float(originalSize[bi].height) / float(netRT->input_dim.h);
 
     // compute dets
     nDets = 0;