专栏首页贾志刚-OpenCV学堂轻松学Pytorch –使用torchvision实现对象检测

轻松学Pytorch –使用torchvision实现对象检测

大家好,前面一篇文章介绍了torchvision的模型ResNet50实现图像分类,这里再给大家介绍一下如何使用torchvision自带的对象检测模型Faster-RCNN实现对象检测。Torchvision自带的对象检测模型是基于COCO数据集训练的,最小分辨率支持800, 最大支持1333的输入图像。

Faster-RCNN模型

Faster-RCNN模型的基础网络是ResNet50, ROI生成使用了RPN,加上头部组成。图示如下:

在torchvision框架下可以通过下面的代码直接下载预训练模型,

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

对模型使用GPU加速支持

# 使用GPU

train_on_gpu = torch.cuda.is_available()
if train_on_gpu:
     model.cuda()

推理输出有三个信息分别为:

boxes:表示对象框
scores:表示每个对象得分
labels:表示对于的分类标签

图像检测

使用模型实现图像检测,支持90个类别的对象检测,代码实现如下:

def faster_rcnn_image_detection():

image = cv.imread("D:/images/cars.jpg")

blob = transform(image)

c, h, w = blob.shape

input_x = blob.view(1, c, h, w)

output = model(input_x.cuda())[0]

boxes = output['boxes'].cpu().detach().numpy()

scores = output['scores'].cpu().detach().numpy()

labels = output['labels'].cpu().detach().numpy()

index = 0

for x1, y1, x2, y2 in boxes:

if scores[index] > 0.5:

print(x1, y1, x2, y2)

cv.rectangle(image, (np.int32(x1), np.int32(y1)),

(np.int32(x2), np.int32(y2)), (0, 255, 255), 1, 8, 0)

label_id = labels[index]

label_txt = coco_names[str(label_id)]

cv.putText(image, label_txt, (np.int32(x1), np.int32(y1)), cv.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 255), 1)

index += 1

cv.imshow("Faster-RCNN Detection Demo", image)

cv.waitKey(0)

cv.destroyAllWindows()

运行结果下:

视频实时对象检测

基于OpenCV实现视频文件或者摄像头读取,完成视频的实时对象检测,代码实现如下:

 1capture = cv.VideoCapture("D:/images/video/vehicle.ts")
 2while True:
 3    ret, frame = capture.read()
 4    if ret is not True:
 5        break
 6    blob = transform(frame)
 7    c, h, w = blob.shape
 8    input_x = blob.view(1, c, h, w)
 9    output = model(input_x.cuda())[0]
10    boxes = output['boxes'].cpu().detach().numpy()
11    scores = output['scores'].cpu().detach().numpy()
12    labels = output['labels'].cpu().detach().numpy()
13    index = 0
14    for x1, y1, x2, y2 in boxes:
15        if scores[index] > 0.5:
16            cv.rectangle(frame, (np.int32(x1), np.int32(y1)),
17                         (np.int32(x2), np.int32(y2)), (0, 255, 255), 1, 8, 0)
18            label_id = labels[index]
19            label_txt = coco_names[str(label_id)]
20            cv.putText(frame, label_txt, (np.int32(x1), np.int32(y1)), cv.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 255), 1)
21        index += 1
22    wk = cv.waitKey(1)
23    if wk == 27:
24        break
25    cv.imshow("video detection Demo", frame)

运行结果如下:

本文分享自微信公众号 - OpenCV学堂(CVSCHOOL),作者:gloomyfish

原文出处及转载信息见文内详细说明,如有侵权,请联系 yunjia_community@tencent.com 删除。

原始发表时间:2020-07-21

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

我来说两句

0 条评论
登录 后参与评论

相关文章

  • 教程 | OpenCV Grabcut对象分割

    Grabcut是基于图割(graph cut)实现的图像分割算法,它需要用户输入一个bounding box作为分割目标位置,实现对目标与背景的分离/分割,这个...

    OpenCV学堂
  • OpenCV中图像频率域滤波

    图像处理不仅可以在空间域进行还可以在频率域进行,把空间域的图像开窗卷积形式,变换得到频率域的矩阵点乘形式得到比较好的效果。转换到频率域最常见的是通过傅里叶变换得...

    OpenCV学堂
  • OpenCV直线拟合检测

    OpenCV直线拟合检测 霍夫直线检测容易受到线段形状与噪声的干扰而失真,这个时候我们需要另辟蹊径,通过对图像进行二值分析,提取骨架,对骨架像素点拟合生成直线,...

    OpenCV学堂
  • 小程序在获取当前位置信息在地图上显示

    小程序在获取当前位置信息在地图上显示api:https://developers.weixin.qq.com/miniprogram/dev/api/wx.ge...

    祈澈菇凉
  • 安全系列之网页防篡改系统

    3手里的一个云项目交付了半年,客户业务系统也顺利通过了云上等保测评二级,为了提前为等保三级做好准备,客户拟购买一些云上的安全服务,需要服务商提供安全咨询,为此我...

    希望的田野
  • 前端学习笔记之Z-index详解

    CSS当中的z-index属性看起来足够简单,但是如果你真的想了解它是如何工作的话,在这简单的表面之下,又有许多值得探究的内容。 在这篇教程中,通过探究层叠上下...

    Jetpropelledsnake21
  • 关于 z-index,你可能一直存在误区

    z-index 这个属性表面看上去很简单,但如果你想搞清楚其工作原理的话,其实是有不少值得探讨之处的。本文将从层叠上下文(stacking contexts)和...

    Chor
  • CSS 层叠上下文(Stacking Context)

    在网页制作的过程中,元素与元素之间的位置关系,在坐标轴上一般可体现为 X 轴、Y 轴和 Z 轴。对于 X 轴和 Y 轴的定位大多数开发都能比较直观的搞清楚,而 ...

    政采云前端团队
  • Python三维绘图--Matplotl

    创建Axes3D主要有两种方式,一种是利用关键字projection='3d'l来实现,另一种则是通过从mpl_toolkits.mplot3d导入对象Axes...

    py3study
  • 007.python科学计算库matplotlib(下)

    版权声明:本文为博主原创文章,允许转载,请标明出处。 https://blog.csdn.net/qwdafedv/article/deta...

    qubianzhong

扫码关注云+社区

领取腾讯云代金券