Skip to content

Commit 2fef5f0

Browse files
authored
yolov5: read input images only once (wang-xinyu#776)
you only read once。 把预处理图片缓存下来,画框时不必读一次,一张图片后处理时间少10ms。
1 parent e785f51 commit 2fef5f0

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

yolov5/yolov5.cpp

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -379,14 +379,16 @@ int main(int argc, char** argv) {
379379
// prepare input data cache in device memory
380380
CUDA_CHECK(cudaMalloc((void**)&img_device, MAX_IMAGE_INPUT_SIZE_THRESH * 3));
381381
int fcount = 0;
382+
std::vector<cv::Mat> imgs_buffer(BATCH_SIZE);
382383
for (int f = 0; f < (int)file_names.size(); f++) {
383384
fcount++;
384385
if (fcount < BATCH_SIZE && f + 1 != (int)file_names.size()) continue;
385386
//auto start = std::chrono::system_clock::now();
386387
float* buffer_idx = (float*)buffers[inputIndex];
387388
for (int b = 0; b < fcount; b++) {
388389
cv::Mat img = cv::imread(img_dir + "/" + file_names[f - fcount + 1 + b]);
389-
if (img.empty()) continue;
390+
if (img.empty()) continue;
391+
imgs_buffer[b] = img;
390392
size_t size_image = img.cols * img.rows * 3;
391393
size_t size_image_dst = INPUT_H * INPUT_W * 3;
392394
//copy data to pinned memory
@@ -408,7 +410,7 @@ int main(int argc, char** argv) {
408410
}
409411
for (int b = 0; b < fcount; b++) {
410412
auto& res = batch_res[b];
411-
cv::Mat img = cv::imread(img_dir + "/" + file_names[f - fcount + 1 + b]);
413+
cv::Mat img = imgs_buffer[b];
412414
for (size_t j = 0; j < res.size(); j++) {
413415
cv::Rect r = get_rect(img, res[j].bbox);
414416
cv::rectangle(img, r, cv::Scalar(0x27, 0xC1, 0x36), 2);

0 commit comments

Comments
 (0)