【pytorch目标检测】实现简单的目标检测(opencv+pytorch)

【pytorch目标检测】实现简单的目标检测(opencv+pytorch),第1张

【pytorch目标检测】实现简单的目标检测(opencv+pytorch) 这次用的是pytorch自带的模型之一:fasterrcnn_resnet50_fpn 因为网络计算量大,所以不容易实现实时的目标检测。

首先,我们先加载一下模型:

model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)

注意的pretrained值true和false的区别,一般我们设置成True就可以,他会在我们的环境里面自动搜索相应的模型参数,如果没有就自动下载

但是因为下载速度太慢,我们就直接使用迅雷下载:

 下载完了之后我们就直接使用model就可以了,毕竟不用咱训练,简单点。。

注意里面的数据增强和归一化,可以直接对opencv对象使用,使其同时成为一个tensor。

model.eval()

cap = cv2.VideoCapture(0)
transform = transforms.Compose(
    [
        transforms.ToTensor()
    ]
)
while True:
    ret, frame = cap.read()
    image = frame
    frame = transform(frame)
    pred = model([frame])

    # 检测出目标的类别和得分
    pred_class = [COCO_INSTANCE_CATEGORY_NAMES[ii] for ii in list(pred[0]['labels'].numpy())]
    pred_score = list(pred[0]['scores'].detach().numpy())

    # 检测出目标的边界框
    pred_boxes = [[ii[0], ii[1], ii[2], ii[3]] for ii in list(pred[0]['boxes'].detach().numpy())]

    # 只保留识别的概率大约 0.5 的结果。
    pred_index = [pred_score.index(x) for x in pred_score if x > 0.5]

    for index in pred_index:
        box = pred_boxes[index]
        cv2.rectangle(img=image, pt1=[int(box[0]), int(box[1])], pt2=[int(box[2]), int(box[3])],
                      color=(0, 0, 225), thickness=3)
        texts = pred_class[index] + ":" + str(np.round(pred_score[index], 2))
        font = cv2.FONT_HERSHEY_SIMPLEX
        cv2.putText(image, texts, (int(box[0]), int(box[1])), font, 1, (200, 255, 155), 2, cv2.LINE_AA)

    cv2.imshow('摄像头', image)
    cv2.waitKey(10)

这里想实现实时检测,但是电脑跑不动。。一卡一卡的。。

 

注意文字的位置,有可能他会加载在图片的外面。。。 

【他能识别出那个床是我没有想到的。。。】

欢迎分享,转载请注明来源:内存溢出

原文地址: https://outofmemory.cn/zaji/5721011.html

(0)
打赏 微信扫一扫 微信扫一扫 支付宝扫一扫 支付宝扫一扫
上一篇 2022-12-18
下一篇 2022-12-18

发表评论

登录后才能评论

评论列表(0条)

保存