在 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 开始到 here 结束的行来加载模型

第 2 步：使用此 function 进行评估。而不是 cv2.imread，你只需要发送你的 frame

第 3 步：按照此function 获取边界框。尤其是这个line。只需引用 't' 变量，您将获得边界框。

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