Pytorch Hub兩行代碼搞定YOLOv5推理

Pytorch Hub介紹

Pytorch Hub是一個幫助研究者實(shí)現(xiàn)模型再現(xiàn)、快速推理驗(yàn)證的預(yù)訓(xùn)練模型庫與一套相關(guān)的API框架。支持遠(yuǎn)程從github上下載指定模型、上傳與分享訓(xùn)練好的模型、支持從本地加載預(yù)訓(xùn)練模型、自定義模型。支持模型遠(yuǎn)程加載與本地推理、當(dāng)前Pytorch Hub已經(jīng)對接到Torchvision、YOLOv5、YOLOv8、pytorchvideo等視覺框架

支持一鍵遠(yuǎn)程加載各種模型，主要包括

支持根據(jù)URL下載指定模型到本地文件夾

YOLOv5支持

首先需要安裝下面的依賴包支持

pip install -r https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt

支持多張圖像推理：

支持本地自定義對象檢測模型推理：

支持多個GPU推理模式

不同設(shè)備之間切換支持

推理參數(shù)支持設(shè)置

相關(guān)源碼貼圖如下：

importtorch
importcv2ascv

#loadimagedata
img="data/images/zidane.jpg"

#加載本地預(yù)訓(xùn)練模型
model=torch.hub.load(repo_or_dir='D:/python/yolov5-7.0-attention/',model='yolov5s',source='local')

#loadimage
#img="D:/bird_test/test004.png"

#加載本地自定義模型
#model=torch.hub.load('D:/python/yolov5-7.0-attention/','custom',path='uav_bird.pt',source='local')
results=model(img)

#顯示
frame=results.render()[0]
bgr=cv.cvtColor(frame,cv.COLOR_RGB2BGR)
cv.imshow("PytorchHub+YOLOv5CustomObjectDetection",bgr)
cv.waitKey(0)

審核編輯：湯梓紅