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达到这个效果。