Yolov5 batch 推理

前言

想要就有了

代码

import shutil
import time
import traceback
import torch
 
import 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 = confidence
    def start(self, img) -> list:
        mode_result = self.model(img,size=640)
        assert mode_result
        detect_result = mode_result.pandas().xyxy[0].to_dict(orient='index')
        if (len(detect_result)) > 0:
            box_list = self._get_box(detect_result)
            return box_list
        else:
            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_list
 
detect=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