在 python 脚本中部署 pytorch .pth 模型

Question

在使用自定义数据集成功训练我的 yolact 模型后，我对 eval.py 使用来自 anaconda 终端的此命令输出的推理结果感到满意：

python eval.py --trained_model=./weights/yolact_plus_resnet50_abrasion_39_10000.pth --config=yolact_resnet_abrasion_config --score_threshold=0.8 --top_k=15 --images=./images:output_images

现在我想运行从我自己的 python 脚本中进行推断，而不是使用 anaconda 终端。我希望能够通过下面的代码获得在网络摄像头帧上进行检测的边界框。有什么想法吗？

import cv2

src = cv2.VideoCapture(0)

while True:
    ret, frame = src.read()
    cv2.imshow('frame', frame)
    key = cv2.waitKey(5)
    if key == (27):
        break

eval.py 代码位于 Yolact 存储库 https://github.com/dbolya/yolact/blob/master/eval.py

Answer 1

我这里只写伪代码给你。

第 1 步：尝试使用从 here and ending here

开始的行加载模型

第 2 步：使用此 function 进行评估。您只需要发送 frame

而不是 cv2.imread

第 3 步：遵循此 function to get the bounding boxes. Especially this line。只需引用 't' 变量，您就会得到边界框。

希望对您有所帮助。如果您需要更多说明，请告诉我。

在 python 脚本中部署 pytorch .pth 模型

Deploy pytorch .pth model in a python script

python

opencv

computer-vision

pytorch