自定义yolov5-v7的推理方法

Posted 2023-03-03 G森森

自定义yolov5-v7的detect.py

修改detect.py代码

import cv2
import torch
import time

def Detect(imgpath,model):
    img=cv2.imread(imgpath)
    img_cvt = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
    results = model(img_cvt)
    results_ = results.pandas().xyxy[0].to_numpy()
    for box in results_:
        # 检测坐标
        l,t,r,b = box[:4].astype('int')
        # 置信度
        confidence = str(round(box[4],2))
        # 类型名称
        name = box[6]
        cv2.rectangle(img,(l,t),(r,b),(0,255,0),2)
        cv2.putText(img,name + " " + confidence,(l,t),cv2.FONT_ITALIC,1,(255,0,0),2)
    cv2.imwrite('bus.jpg',img)

if __name__ == "__main__":
    model = torch.hub.load('', 'custom', path='yolov5s.pt', source='local')
    model.conf = 0.4
    start_time = time.time()
    Detect('./data/images/bus.jpg',model)
    end_time = time.time()  # 程序结束时间
    run_time = end_time - start_time  # 程序的运行时间，单位为秒
    print("运行时间：" + str(run_time) + "秒")

小问题

突然发现这边调用的torch.hub.load（）方法比较慢，自己虽然电脑不好，但用了6秒左右也太夸张了。
目前我采用启动程序后直接生成那个model对象，等待外部触发直接调用这个model就可以节约时间了。但我想是否能节约这方面的时间，欢迎大家给予建议，谢谢！