当前位置：首页 > 编程日记 > 正文

TorchVision中使用FasterRCNN+ResNet50+FPN进行目标检测

编程日记 2024-12-13 13:00:00

TorchVision中给出了使用ResNet-50-FPN主干(backbone)构建Faster R-CNN的pretrained模型，模型存放位置为https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth，可通过fasterrcnn_resnet50_fpn函数下载，此函数实现在torchvison/models/detection/faster_rcnn.py中，下载后在Ubuntu上存放在~/.cache/torch/hub/checkpoints目录下，在Windows上存放在C:\Users\spring\.cache\torch\hub\checkpoints目录下，其中spring为用户名。

模型的输入是一个tensor列表；每个shape都是[c,h,w]；每个shape指定一副图像，并且图像中值的范围为[0,1]，即已做过normalized；不同的图像可以有不同的大小，即此模型支持非固定大小图像的输入。

模型的行为取决于它是处于训练模式(training)还是评估模式(evaluation)：

(1).在训练期间，模型需要输入tensors和targets(字典列表)，包含boxes和labels。

boxes类型为FloatTensor[N,4]，其中N为图像数；4为[x1,y1,x2,y2]，即ground-truth box的左上和右下角坐标，它们的值要合理范围内。

labels类型为Int64Tensor[N]，每个ground-truth box的class label。

(2).在推理(inference)过程中，模型只需要输入tensors，并返回后处理的预测(post-processed predictions)，此预测类型为List[Dict[Tensor]]，对应每个输入图像。

Dict字段内容除包含boxes和labels外，还包含scores。

scores类型为Tensor[N]，每个预测的分值，按照值从大到小的顺序排列。

模型是通过COCO数据集训练获得的，COCO数据集的介绍参考：https://blog.csdn.net/fengbingchun/article/details/121308708

FPN全称为Feature Pyramid Networks，即特征金字塔网络，是一种多尺度的目标检测算法，FPN的介绍参考：https://blog.csdn.net/fengbingchun/article/details/87359191

ResNet即Residual Networks，也称为残差网络，是为了解决深度神经网络的”退化(degradation)”问题。ResNet-50中的50指此网络有50层。ResNet介绍参考：https://blog.csdn.net/fengbingchun/article/details/114167581

Faster R-CNN为目标检测算法，为RPN(Region Proposal Network)和Fast R-CNN的结合。Faster R-CNN介绍参考：https://blog.csdn.net/fengbingchun/article/details/87195597

以下为测试代码：

import torch
from torchvision import models
from torchvision import transforms
import cv2'''
Note: conda pytorch install opencv
windows: conda install opencv # python=3.8.8, opencv=4.0.1
ubuntu: pip3 install opencv-python # python=3.7.11, opencv=4.5.4
'''images_path = "../../data/image/"
images_name = ["1.jpg", "2.jpg", "4.jpg"]
images_data = [] # opencv
tensor_data = [] # pytorch tensorfor name in images_name:img = cv2.imread(images_path + name)print(f"name: {images_path+name}, opencv image shape: {img.shape}") # (w,h,c)images_data.append(img)transform = transforms.Compose([transforms.ToTensor()])tensor = transform(img) # Normalized Tensor image: [0., 1.]print(f"tensor shape: {tensor.shape}, max: {torch.max(tensor)}, min: {torch.min(tensor)}") # (c,h,w)tensor_data.append(tensor)# reference: torchvison/models/detection/faster_rcnn.py
# 使用ResNet-50-FPN(Feature Pyramid Networks, 特征金字塔网络)构建Faster RCNN模型
model = models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
#print(model) # 可查看模型结构
model.eval() # 推理
predictions = model(tensor_data) # result: list: boxes (FloatTensor[N, 4]), labels (Int64Tensor[N]), scores (Tensor[N])
#print(predictions)coco_labels_name = ["unlabeled", "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat","traffic light", "fire hydrant", "street sign", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse","sheep", "cow", "elephant", "bear", "zebra", "giraffe", "hat", "backpack", "umbrella", "shoe","eye glasses", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports_ball", "kite", "baseball bat","baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "plate", "wine glass", "cup", "fork", "knife","spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot_dog", "pizza","donut", "cake", "chair", "couch", "potted plant", "bed", "mirror", "dining table", "window", "desk","toilet", "door", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven","toaster", "sink", "refrigerator", "blender", "book", "clock", "vase", "scissors", "teddy bear", "hair drier","toothbrush", "hair brush"] # len = 92for x in range(len(predictions)):pred = predictions[x]scores = pred["scores"]mask = scores > 0.5 # 只取scores值大于0.5的部分boxes = pred["boxes"][mask].int().detach().numpy() # [x1, y1, x2, y2]labels = pred["labels"][mask]scores = scores[mask]print(f"prediction: boxes:{boxes}, labels:{labels}, scores:{scores}")img = images_data[x]for idx in range(len(boxes)):cv2.rectangle(img, (boxes[idx][0], boxes[idx][1]), (boxes[idx][2], boxes[idx][3]), (255, 0, 0))cv2.putText(img, coco_labels_name[labels[idx]]+" "+str(scores[idx].detach().numpy()), (boxes[idx][0]+10, boxes[idx][1]+10), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 0), 1)cv2.imshow("image", img)cv2.waitKey(1000)cv2.imwrite(images_path+"result_"+images_name[x], img)print("test finish")

说明：

(1).输入图像既可以是彩色图也可以是灰度图，即channel为3或1均可。

(2).输入图像的大小不受限制，一组图像可以大小不一致。

(3).输入图像要求normalized到[0., 1.]。

(4).执行结果仅显示scores值大于0.5的情况。

(5).测试代码中类别数为92而不是80，92=1+11+80。其中1为id为0，label name为unlabeled；11为从COCO中移除的label，如street sign；80为真正的label数，如person。详细参考：https://github.com/nightrome/cocostuff/blob/master/labels.md

(6).结果显示中有冗余的检测框，可以通过NMS(Non-Maximum Suppression)非极大值抑制算法移除。

执行结果如下：以下原始测试图像来自网络

GitHub： GitHub - fengbingchun/PyTorch_Test: PyTorch's usage

https://www.dkcj.cn/info/31648.html

TorchVision中使用FasterRCNN+ResNet50+FPN进行目标检测

相关文章：

iOS-UIButton防止重复点击(三种办法)

华为诺亚方舟开源预训练模型“哪吒”，4项任务均达到SOTA

音量调节助手(转)

TorchVision中通过AlexNet网络进行图像分类

西湖龙井也上链？是的，以后你喝什么茶我都知道！

Java Enumeration接口

Windows Azure Pack与SCVMM标签解析分享

Linux下C++中可使用的3种Hook方法

Java Properties 类

国产数据库年终大盘点

解密FFmpeg播放track mode控制

UIButton防止按钮和手势的暴力点击

使用Windows7上的VS Code打开远程机Ubuntu上的文件操作步骤

微众银行殷磊：AI+卫星，从上帝视角洞察资产管理｜BDTC 2019

【二分答案】【最短路】bzoj1614 [Usaco2007 Jan]Telephone Lines架设电话线

Python3中装饰器@typing.overload的使用

19年NAACL纪实：自然语言处理的实用性见解 | CSDN博文精选

高并发场景下数据库的常见问题及解决方案

典型用户 persona

PyTorch中nn.Module类简介

什么是三层交换机、网关、DNS、子网掩码、MAC地址

20行代码发一篇NeurIPS：梯度共享已经不安全了

关于页面打印window.print()的样式问题

Python3中参数*args和**kwargs介绍

4大主流CPU处理器技术架构，不知道就out了！

grunt-connect-proxy解决开发时跨域问题

混合推荐系统就是多个推荐系统“大杂烩”吗？

Python3中collections.OrderedDict介绍

汗！雅虎中国个人空间

关于v$process与v$session中process的理解