Yolov5 batch 推理

前言

想要就有了

代码

import shutil
import time
import traceback
import torchimport os
import cv2
class PeopleDetect(object):def __init__(self, repo_or_dir, weight_path, confidence) -> None:self.model = torch.hub.load(repo_or_dir, "custom", path=weight_path, source='local',force_reload=True)self.confidence = confidencedef start(self, img) -> list:mode_result = self.model(img,size=640)assert mode_resultdetect_result = mode_result.pandas().xyxy[0].to_dict(orient='index')if (len(detect_result)) > 0:box_list = self._get_box(detect_result)return box_listelse:return []def _get_box(self, detect_result:list) -> list:box_list = []for i in range(len(detect_result)):conf = detect_result[i]['confidence']xmin = int(detect_result[i]['xmin'])ymin = int(detect_result[i]['ymin'])xmax = int(detect_result[i]['xmax'])ymax = int(detect_result[i]['ymax'])if conf > self.confidence:box_list.append([xmin, ymin, xmax, ymax,conf])return box_listdetect=PeopleDetect(repo_or_dir="yolov5", weight_path='weight/best.pt', confidence=0.01)images_path="/home/save_images"
save_path="/home/errors"
Filelist = []
for home, dirs, files in os.walk(images_path):for filename in files:Filelist.append(os.path.join(home, filename))start_time=time.time()
image_list=[]
count=0
for ind,i in enumerate( Filelist):try:image_ori=cv2.imread(i)if len(image_list)<128:image_list.append(image_ori[:,:,-1])else:res=detect.start(image_list)image_list=[]print(int(time.time() - start_time) / int(ind))except (Exception,BaseException)as e:print(traceback.format_exc())

测试结果

以下结果使用RTX3080 测试得出，imgsize为640，model为5s