Skip to content

Commit 8d32c30

Browse files
Nuzhny007Nuzhny007
committed
Batch inference for TensorRT
1 parent b18055c commit 8d32c30

19 files changed

Lines changed: 204 additions & 123 deletions

‎example/examples.h‎

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -645,9 +645,10 @@ class YoloDarknetExample : public VideoExample
645645
{
646646
TinyYOLOv3 = 0,
647647
YOLOv3,
648-
YOLOv4
648+
YOLOv4,
649+
TinyYOLOv4
649650
};
650-
YOLOModels usedModel = YOLOModels::YOLOv4;
651+
YOLOModels usedModel = YOLOModels::YOLOv4;
651652
switch (usedModel)
652653
{
653654
case YOLOModels::TinyYOLOv3:
@@ -667,9 +668,15 @@ class YoloDarknetExample : public VideoExample
667668
config.emplace("modelBinary", pathToModel + "yolov4.weights");
668669
config.emplace("confidenceThreshold", "0.5");
669670
break;
671+
672+
case YOLOModels::TinyYOLOv4:
673+
config.emplace("modelConfiguration", pathToModel + "yolov4-tiny.cfg");
674+
config.emplace("modelBinary", pathToModel + "yolov4-tiny.weights");
675+
config.emplace("confidenceThreshold", "0.5");
676+
break;
670677
}
671678
config.emplace("classNames", pathToModel + "coco.names");
672-
config.emplace("maxCropRatio", "-1");
679+
config.emplace("maxCropRatio", "1");
673680

674681
config.emplace("white_list", "person");
675682
config.emplace("white_list", "car");
@@ -836,6 +843,7 @@ class YoloTensorRTExample : public VideoExample
836843
config.emplace("confidenceThreshold", "0.5");
837844
config.emplace("inference_precison", "FP32");
838845
config.emplace("net_type", "YOLOV3_TINY");
846+
config.emplace("maxBatch", "4");
839847
break;
840848

841849
case YOLOModels::YOLOv3:
@@ -844,14 +852,16 @@ class YoloTensorRTExample : public VideoExample
844852
config.emplace("confidenceThreshold", "0.7");
845853
config.emplace("inference_precison", "FP32");
846854
config.emplace("net_type", "YOLOV3");
855+
config.emplace("maxBatch", "2");
847856
break;
848857

849858
case YOLOModels::YOLOv4:
850859
config.emplace("modelConfiguration", pathToModel + "yolov4.cfg");
851860
config.emplace("modelBinary", pathToModel + "yolov4.weights");
852-
config.emplace("confidenceThreshold", "0.5");
861+
config.emplace("confidenceThreshold", "0.8");
853862
config.emplace("inference_precison", "FP32");
854863
config.emplace("net_type", "YOLOV4");
864+
config.emplace("maxBatch", "1");
855865
break;
856866

857867
case YOLOModels::TinyYOLOv4:
@@ -860,11 +870,12 @@ class YoloTensorRTExample : public VideoExample
860870
config.emplace("confidenceThreshold", "0.5");
861871
config.emplace("inference_precison", "FP32");
862872
config.emplace("net_type", "YOLOV4_TINY");
873+
config.emplace("maxBatch", "4");
863874
break;
864875
}
865876

866877
config.emplace("classNames", pathToModel + "coco.names");
867-
config.emplace("maxCropRatio", "-1");
878+
config.emplace("maxCropRatio", "2");
868879

869880
config.emplace("white_list", "person");
870881
config.emplace("white_list", "car");

‎src/Detector/BaseDetector.h‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,57 @@ class BaseDetector
120120
cv::Size m_minObjectSize;
121121

122122
cv::Mat m_motionMap;
123+
124+
std::vector<cv::Rect> GetCrops(float maxCropRatio, cv::Size netSize, cv::Size imgSize) const
125+
{
126+
std::vector<cv::Rect> crops;
127+
128+
const float whRatio = static_cast<float>(netSize.width) / static_cast<float>(netSize.height);
129+
int cropHeight = cvRound(maxCropRatio * netSize.height);
130+
int cropWidth = cvRound(maxCropRatio * netSize.width);
131+
132+
if (imgSize.width / (float)imgSize.height > whRatio)
133+
{
134+
if (cropHeight >= imgSize.height)
135+
cropHeight = imgSize.height;
136+
cropWidth = cvRound(cropHeight * whRatio);
137+
}
138+
else
139+
{
140+
if (cropWidth >= imgSize.width)
141+
cropWidth = imgSize.width;
142+
cropHeight = cvRound(cropWidth / whRatio);
143+
}
144+
145+
//std::cout << "Frame size " << imgSize << ", crop size = " << cv::Size(cropWidth, cropHeight) << ", ratio = " << maxCropRatio << std::endl;
146+
147+
const int stepX = 3 * cropWidth / 4;
148+
const int stepY = 3 * cropHeight / 4;
149+
for (int y = 0; y < imgSize.height; y += stepY)
150+
{
151+
bool needBreakY = false;
152+
if (y + cropHeight >= imgSize.height)
153+
{
154+
y = imgSize.height - cropHeight;
155+
needBreakY = true;
156+
}
157+
for (int x = 0; x < imgSize.width; x += stepX)
158+
{
159+
bool needBreakX = false;
160+
if (x + cropWidth >= imgSize.width)
161+
{
162+
x = imgSize.width - cropWidth;
163+
needBreakX = true;
164+
}
165+
crops.emplace_back(x, y, cropWidth, cropHeight);
166+
if (needBreakX)
167+
break;
168+
}
169+
if (needBreakY)
170+
break;
171+
}
172+
return crops;
173+
};
123174
};
124175

125176

‎src/Detector/OCVDNNDetector.cpp‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -244,12 +244,12 @@ void OCVDNNDetector::Detect(cv::UMat& colorFrame)
244244
}
245245

246246
std::cout << "cropsCount = " << cropsCount << std::endl;
247-
if (cropsCount > 1)
248-
nms3<CRegion>(tmpRegions, m_regions, m_nmsThreshold,
249-
[](const CRegion& reg) -> cv::Rect { return reg.m_brect; },
250-
[](const CRegion& reg) -> float { return reg.m_confidence; },
251-
[](const CRegion& reg) -> std::string { return reg.m_type; },
252-
0, 0.f);
247+
if (cropsCount > 1)
248+
nms3<CRegion>(tmpRegions, m_regions, m_nmsThreshold,
249+
[](const CRegion& reg) { return reg.m_brect; },
250+
[](const CRegion& reg) { return reg.m_confidence; },
251+
[](const CRegion& reg) { return reg.m_type; },
252+
0, 0.f);
253253
}
254254
}
255255

‎src/Detector/SSDMobileNetDetector.cpp‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -183,10 +183,10 @@ void SSDMobileNetDetector::Detect(cv::UMat& colorFrame)
183183
}
184184

185185
nms3<CRegion>(tmpRegions, m_regions, 0.4f,
186-
[](const CRegion& reg) -> cv::Rect { return reg.m_brect; },
187-
[](const CRegion& reg) -> float { return reg.m_confidence; },
188-
[](const CRegion& reg) -> std::string { return reg.m_type; },
189-
0, 0.f);
186+
[](const CRegion& reg) { return reg.m_brect; },
187+
[](const CRegion& reg) { return reg.m_confidence; },
188+
[](const CRegion& reg) { return reg.m_type; },
189+
0, 0.f);
190190
}
191191

192192
///

‎src/Detector/YoloDarknetDetector.cpp‎

Lines changed: 11 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,6 @@ bool YoloDarknetDetector::Init(const config_t& config)
4646

4747
m_detector = std::make_unique<Detector>(modelConfiguration->second, modelBinary->second, currGPUID);
4848
m_detector->nms = 0.2f;
49-
m_WHRatio = static_cast<float>(m_detector->get_net_width()) / static_cast<float>(m_detector->get_net_height());
5049

5150
auto classNames = config.find("classNames");
5251
if (classNames != config.end())
@@ -96,61 +95,20 @@ void YoloDarknetDetector::Detect(cv::UMat& colorFrame)
9695
}
9796
else
9897
{
99-
int cropHeight = cvRound(m_maxCropRatio * m_detector->get_net_height());
100-
int cropWidth = cvRound(m_maxCropRatio * m_detector->get_net_width());
101-
102-
if (colorFrame.cols / (float)colorFrame.rows > m_WHRatio)
103-
{
104-
if (cropHeight >= colorFrame.rows)
105-
cropHeight = colorFrame.rows;
106-
cropWidth = cvRound(cropHeight * m_WHRatio);
107-
}
108-
else
109-
{
110-
if (cropWidth >= colorFrame.cols)
111-
cropWidth = colorFrame.cols;
112-
cropHeight = cvRound(colorFrame.cols / m_WHRatio);
113-
}
114-
115-
//std::cout << "Frame size " << colorFrame.size() << ", crop size = " << cv::Size(cropWidth, cropHeight) << ", ratio = " << m_maxCropRatio << std::endl;
116-
117-
cv::Rect crop(0, 0, cropWidth, cropHeight);
118-
regions_t tmpRegions;
119-
size_t cropsCount = 0;
120-
int stepX = 3 * crop.width / 4;
121-
int stepY = 3 * crop.height / 4;
122-
for (; crop.y < colorMat.rows; crop.y += stepY)
123-
{
124-
bool needBreakY = false;
125-
if (crop.y + crop.height >= colorMat.rows)
126-
{
127-
crop.y = colorMat.rows - crop.height;
128-
needBreakY = true;
129-
}
130-
for (crop.x = 0; crop.x < colorMat.cols; crop.x += stepX)
131-
{
132-
bool needBreakX = false;
133-
if (crop.x + crop.width >= colorMat.cols)
134-
{
135-
crop.x = colorMat.cols - crop.width;
136-
needBreakX = true;
137-
}
138-
139-
std::cout << "Crop " << cropsCount++ << ": " << crop << std::endl;
140-
DetectInCrop(colorMat, crop, tmpRegions);
141-
142-
if (needBreakX)
143-
break;
144-
}
145-
if (needBreakY)
146-
break;
147-
}
98+
std::vector<cv::Rect> crops = GetCrops(m_maxCropRatio, cv::Size(m_detector->get_net_width(), m_detector->get_net_height()), colorMat.size());
99+
regions_t tmpRegions;
100+
for (size_t i = 0; i < crops.size(); ++i)
101+
{
102+
const auto& crop = crops[i];
103+
std::cout << "Crop " << i << ": " << crop << std::endl;
104+
DetectInCrop(colorMat, crop, tmpRegions);
105+
}
148106

149107
//std::cout << "nms for " << tmpRegions.size() << " objects" << std::endl;
150108
nms3<CRegion>(tmpRegions, m_regions, 0.4f,
151-
[](const CRegion& reg) -> cv::Rect { return reg.m_brect; },
152-
[](const CRegion& reg) -> float { return reg.m_confidence; },
153-
[](const CRegion& reg) -> std::string { return reg.m_type; },
109+
[](const CRegion& reg) { return reg.m_brect; },
110+
[](const CRegion& reg) { return reg.m_confidence; },
111+
[](const CRegion& reg) { return reg.m_type; },
154112
0, 0.f);
155113
}
156114
//std::cout << "Finally " << m_regions.size() << " objects" << std::endl;

‎src/Detector/YoloDarknetDetector.h‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,6 @@ class YoloDarknetDetector : public BaseDetector
3030
private:
3131
std::unique_ptr<Detector> m_detector;
3232

33-
float m_WHRatio = 1.f;
3433
float m_confidenceThreshold = 0.5f;
3534
float m_maxCropRatio = 3.0f;
3635
std::vector<std::string> m_classNames;

‎src/Detector/YoloDetector.cpp‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -224,11 +224,11 @@ void YoloOCVDetector::Detect(cv::UMat& colorFrame)
224224
}
225225
}
226226

227-
nms3<CRegion>(tmpRegions, m_regions, 0.4f,
228-
[](const CRegion& reg) -> cv::Rect { return reg.m_brect; },
229-
[](const CRegion& reg) -> float { return reg.m_confidence; },
230-
[](const CRegion& reg) -> std::string { return reg.m_type; },
231-
0, 0.f);
227+
nms3<CRegion>(tmpRegions, m_regions, 0.4f,
228+
[](const CRegion& reg) { return reg.m_brect; },
229+
[](const CRegion& reg) { return reg.m_confidence; },
230+
[](const CRegion& reg) { return reg.m_type; },
231+
0, 0.f);
232232
}
233233

234234
///

‎src/Detector/YoloTensorRTDetector.cpp‎

Lines changed: 51 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#include <fstream>
22
#include "YoloTensorRTDetector.h"
3+
#include "nms.h"
34

45
///
56
/// \brief YoloTensorRTDetector::YoloTensorRTDetector
@@ -119,15 +120,55 @@ bool YoloTensorRTDetector::Init(const config_t& config)
119120
void YoloTensorRTDetector::Detect(cv::UMat& colorFrame)
120121
{
121122
m_regions.clear();
123+
cv::Mat colorMat = colorFrame.getMat(cv::ACCESS_READ);
124+
125+
if (m_maxCropRatio <= 0)
126+
{
127+
std::vector<cv::Mat> batch = { colorMat };
128+
std::vector<tensor_rt::BatchResult> detects;
129+
m_detector->detect(batch, detects);
130+
for (const tensor_rt::BatchResult& dets : detects)
131+
{
132+
for (const tensor_rt::Result& bbox : dets)
133+
{
134+
m_regions.emplace_back(bbox.rect, m_classNames[bbox.id], bbox.prob);
135+
}
136+
}
137+
}
138+
else
139+
{
140+
std::vector<cv::Rect> crops = GetCrops(m_maxCropRatio, m_detector->get_input_size(), colorMat.size());
141+
regions_t tmpRegions;
142+
for (size_t i = 0; i < crops.size();)
143+
{
144+
size_t batchsize = std::min(static_cast<size_t>(m_localConfig.n_max_batch), crops.size() - i);
145+
std::vector<cv::Mat> batch;
146+
batch.reserve(batchsize);
147+
for (size_t j = 0; j < batchsize; ++j)
148+
{
149+
batch.emplace_back(colorMat, crops[i + j]);
150+
}
151+
std::vector<tensor_rt::BatchResult> detects;
152+
m_detector->detect(batch, detects);
153+
154+
for (size_t j = 0; j < batchsize; ++j)
155+
{
156+
const auto& crop = crops[i + j];
157+
//std::cout << "Crop " << (i + j) << ": " << crop << std::endl;
122158

123-
std::vector<cv::Mat> batch = { colorFrame.getMat(cv::ACCESS_READ) };
124-
std::vector<tensor_rt::BatchResult> detects;
125-
m_detector->detect(batch, detects);
126-
for (const tensor_rt::BatchResult& dets : detects)
127-
{
128-
for (const tensor_rt::Result& bbox : dets)
129-
{
130-
m_regions.emplace_back(bbox.rect, m_classNames[bbox.id], bbox.prob);
131-
}
132-
}
159+
for (const tensor_rt::Result& bbox : detects[j])
160+
{
161+
tmpRegions.emplace_back(cv::Rect(bbox.rect.x + crop.x, bbox.rect.y + crop.y, bbox.rect.width, bbox.rect.height), m_classNames[bbox.id], bbox.prob);
162+
}
163+
}
164+
i += batchsize;
165+
}
166+
167+
nms3<CRegion>(tmpRegions, m_regions, 0.4f,
168+
[](const CRegion& reg) { return reg.m_brect; },
169+
[](const CRegion& reg) { return reg.m_confidence; },
170+
[](const CRegion& reg) { return reg.m_type; },
171+
0, 0.f);
172+
//std::cout << "nms for " << tmpRegions.size() << " objects - result " << m_regions.size() << std::endl;
173+
}
133174
}

‎src/Detector/tensorrt_yolo/calibrator.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ Int8EntropyCalibrator::Int8EntropyCalibrator(const uint32_t& batchSize, const st
5555

5656
Int8EntropyCalibrator::~Int8EntropyCalibrator() { NV_CUDA_CHECK(cudaFree(m_DeviceInput)); }
5757

58-
bool Int8EntropyCalibrator::getBatch(void* bindings[], const char* names[], int nbBindings)
58+
bool Int8EntropyCalibrator::getBatch(void* bindings[], const char* names[], int /*nbBindings*/)
5959
{
6060
if (m_ImageIndex + m_BatchSize >= m_ImageList.size()) return false;
6161

‎src/Detector/tensorrt_yolo/chunk.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ namespace nvinfer1
4444
cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) override;
4545
void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) override;
4646
void detachFromContext() override;
47-
bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const override
47+
bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int /*nbInputs*/, int /*nbOutputs*/) const override
4848
{
4949
return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT;
5050
}

0 commit comments

Comments
 (0)