如何在自定义Python脚本中部署YOLACT .pth模型获取摄像头检测边界框

前置说明 #

你需要将编写的Python脚本放在YOLACT项目的根目录下运行，确保可以正常导入项目内的所有依赖模块。下述实现逻辑完全复用官方eval.py的推理逻辑，参数和你之前在终端执行的命令完全对齐，不需要额外修改适配。

完整实现代码 #

import cv2
import torch
from yolact import Yolact
from data import cfg, set_cfg
from utils.augmentations import BaseTransform
from utils.functions import detect

# 1. 加载配置，和你终端命令的参数对齐
model_path = './weights/yolact_plus_resnet50_abrasion_39_10000.pth'
config_name = 'yolact_resnet_abrasion_config'
score_threshold = 0.8
top_k = 15
set_cfg(config_name)

# 2. 初始化模型
net = Yolact()
net.load_weights(model_path)
net.eval()
net = net.cuda() # 如果你用CPU推理的话把这行删掉，下一行改成net = net.cpu()
transform = BaseTransform(net.size, cfg.dataset.mean, cfg.dataset.std)

# 3. 初始化摄像头
src = cv2.VideoCapture(0)
while True:
    ret, frame = src.read()
    if not ret:
        break
    
    # 4. 帧预处理+推理
    frame_tensor = torch.from_numpy(transform(frame)[0]).permute(2, 0, 1)
    with torch.no_grad():
        preds = net(frame_tensor.unsqueeze(0).cuda()) # CPU推理的话把.cuda()删掉
    dets = detect(preds, frame.shape[0], frame.shape[1], score_threshold=score_threshold, top_k=top_k)[0]
    
    # 5. 提取边界框、置信度、类别ID，dets的格式为[类别数，每个类别的检测结果，5=4个坐标+1置信度]
    bbox_list = []
    for class_id in range(len(dets)):
        for detection in dets[class_id]:
            if detection[4] < score_threshold:
                continue
            # 边界框格式为 [x1, y1, x2, y2]，均为原图的像素坐标
            x1, y1, x2, y2 = detection[:4]
            score = detection[4]
            bbox_list.append({
                'class_id': class_id,
                'bbox': [int(x1), int(y1), int(x2), int(y2)],
                'score': float(score)
            })
            # 可选：在帧上画边界框可视化
            cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0,255,0), 2)
            cv2.putText(frame, f'cls:{class_id} score:{score:.2f}', (int(x1), int(y1)-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 2)
    
    # 输出所有检测到的边界框，你可以在这里加自己的处理逻辑
    if len(bbox_list) > 0:
        print(f'当前帧检测到{len(bbox_list)}个目标，边界框列表：{bbox_list}')

    cv2.imshow('frame', frame)
    key = cv2.waitKey(5)
    if key == 27:
        break

src.release()
cv2.destroyAllWindows()

注意事项 #

• 如果你是CPU环境推理，删除代码中对应的.cuda()后缀，同时修改模型加载的设备配置即可
• 如果你需要获取实例分割掩码，可以在detect函数返回结果中提取对应字段，逻辑和提取边界框一致
• 所有参数和你之前终端执行的命令完全匹配，推理效果和终端运行eval.py的效果完全一致

内容的提问来源于stack exchange，提问作者Ramiro Sena