【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)

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

 

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

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