pytorch 实时物体检测

前言

使用了pytorch自己提供的maskrcnn_resnet50_fpn，也就是说，不过是调用人家训练好的代码而已，当个玩具了，供大家参考。

github项目

其实主要的代码量在结果的显示上，并不是模型的调用上。
这个是github项目链接.
效果：

实现思路

模型调用，这个就不必细说了。

• 使用opencv读图片
• 然后转化成torch能吃进去的数据，
• input到模型中，得到返回值。根据模型的返回值，一一把label标注到原图片中
• 然后显示图片

下面的是主要代码：
不能直接运行，只是思路，完整的项目在前面已经给出。

id2name作用是把模型返回值id直接映射到他的类型。
colors是用彩色标注不同类型的框框和字体，便于查看。

import torchvision
from torchvision import transforms
from utils import *
import torch

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

win_name='detect'

# download model
model=torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)

# read image and transform it
trans=transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0,0,0],[1,1,1])
])


model=model.eval()
model=model.to(device)
cam = cv2.VideoCapture(0)

# 80 classes in total
id2name=GetClass()
# prepare colors for every boxes
colors=np.random.randint(0, 255,(len(id2name),1,3))

with torch.no_grad():
    while True:
        ret, original = cam.read()
        # cv2.imshow(win_name, mat)
        mat = np.array(cv2.cvtColor(original, cv2.COLOR_BGR2RGB))
        frame = trans(mat).unsqueeze(0)
        frame=frame.to(device)
        prediction = model(frame)
        try:
            labels = prediction[0]['labels']
            boxes = prediction[0]['boxes']
            scores = prediction[0]['scores']

            ShowPicResult(original, win_name, scores, labels, boxes,id2name,colors,False)
        except :
            print('something went wrong!')
        if cv2.waitKey(20) & 0xFF == ord('q'):
            break

    cam.release()
    cv2.destroyAllWindows()

一些项目说明

• 毕竟是图像的应用，所以，最好电脑上装了N卡的cuda套件，不然实时不了，因为太慢了。
• maskrcnn本身使用coco的训练集，这个训练集的id和对应种类不是直接对应的，具体的对应关系我在博客中也有说明，具体看这里。
• 这个项目是需要启用摄像头的，没有摄像头的电脑的同学可以改改代码，用上自己的视频。

最后

以上就是壮观皮带为你收集整理的pytorch 实时物体检测的全部内容，希望文章能够帮你解决pytorch 实时物体检测所遇到的程序开发问题。

如果觉得靠谱客网站的内容还不错，欢迎将靠谱客网站推荐给程序员好友。