Yolov5 将图片合成batch进行推理-1

Yolov5 将图片合成batch进行推理-1

代码

import os
import cv2
import torch
import time

s_t = time.time()
path_weight = '/home/yz5/cq/yolov5_24bit/weights/yolov5s.pt'
model = torch.hub.load('/home/yz5/cq/yolov5_24bit', 'custom', path = path_weight, source = 'local', autoshape = True)

img_list = []
dir_path = 'yolov5_24bit/data/PCBDatasets/images'
for i in os.listdir(dir_path):
    img_list.append(cv2.imread(dir_path + '/' + i)[..., ::-1])

results = model(img_list, size=512)  # batch of images

results.print()

print("Cost Time：", time.time() - s_t)

解析

代码中加载的路径'/home/yz5/cq/yolov5_24bit'实际上是Yolo文件夹中的hubconf.py的所在文件夹。这个例子就是将'yolov5_24bit/data/PCBDatasets/images'下的图片全部合成一个batch送到model中进行检测。

但是无论是把20-100张图片合成一个batch送去检测，都比逐张检测是要慢的。

暂时没有弄清楚怎么通过这种方式给model输送多个batch。后续可能根据上面那种模式添加代码或者通过修改detect.py达到这个效果。