PyTorch/[PyTorch 学习笔记] 8.2 目标检测简介
本章代码:
这篇文章主要介绍了目标检测。
目标检测是判断目标在图像中的位置,有两个要素:
- 分类:分类向量\(P_{0}, P_{1}, P_{2}...\),shape 为\([N, c+1]\)
- 回归:回归边界框\([x_{1}, x_{2}, y_{1}, y_{2}]\),shape 为\([n, 4]\)
下面代码是加载预训练好的FasterRCNN_ResNet50_fpn
,这个模型在是 COCO 模型上进行训练的,有 91 种类别。这里图片不再是BCHW
的形状,而是一个list
,每个元素是图片。输出也是一个 list,每个元素是一个 dict,每个 dict 包含三个元素:boxes、scores、labels,每个元素都是 list,因为一张图片中可能包含多个目标。接着是绘制框的代码,scores
的的某个元素小于某个阈值,则不绘制这个框。
1 | import os |
输出如下:
目标检测中难题之一是边界框的数量 \(N\) 的确定。
传统的方法是滑动窗口策略,缺点是重复计算量大,窗口大小难以确定。
将全连接层改为卷积层,最后一层特征图的一个像素就对应着原图的一个区域,就可以使用利用卷积操作实现滑动窗口。
目标检测模型可以划分为 one-stage 和 two-stage。
one-stage 包括:
- YOLO
- SSD
- Retina-Net
two-stage 包括:
- RCNN
- SPPNet
- Fast RCNN
- Faster RCNN
- Pyramid Network
one-stage 的模型是直接把得到的特征图划分为多个网格,每个网格分别做分类和回归。
two-stage 的模型多了 proposal generation,输出 多个候选框,通常默认 2000 个候选框
在 Faster RCNN 中,proposal generation 是 RPN(Region Proposal Network),会根据 feature map 生成数十万个候选框,通过 NMS 选出前景概率最高的 2000 个框。由于候选框的大小各异,通过 ROI pooling,得到固定大小的输出,channel 数量就是框的数量。ROI pooling 的特点是输入特征图尺寸不固定,但是输出特征图尺寸固定。最后经过全连接层得到回归和分类的输出。
fasterrcnn_resnet50_fpn
会返回一个FasterRCNN
,FasterRCNN
继承于GeneralizedRCNN
,GeneralizedRCNN
的forward()
函数中包括下面 3 个模块:
backbone:
features = self.backbone(images.tensors)
在构建 backbone 时,会根据
backbone_name
选择对应的 backbone,这里使用 resnet50。rpn:
proposals, proposal_losses = self.rpn(images, features, targets)
roi_heads:
detections, detector_losses = self.roi_heads(features, proposals, images.image_sizes, targets)
GeneralizedRCNN
的forward()
函数如下:
1 | def forward(self, images, targets=None): |
self.backbone(images.tensors)
返回的features
是一个 dict,每个元素是一个 feature map,每个特征图的宽高是上一个特征图宽高的一半。
这 5 个 feature map 分别对应 ResNet 中的 5 个特征图
接下来进入 rpn 网络,rpn 网络代码如下。
1 | def forward(self, images, features, targets=None): |
self.head(features)
会调用RPNHead
,返回的objectness
和pred_bbox_deltas
都是和features
大小一样的 dict,只是 channel 数量不一样。objectness
的 channel 数量是 3,表示特征图的一个像素点输出 3 个可能的框;pred_bbox_deltas
的 channel 数量是 12,表示每个特征图的 3 个框的坐标的偏移量。
self.anchor_generator(images, features)
的输出是anchors
,形状是\(242991 \times 4\),这是真正的框。
self.filter_proposals()
对应的是 NMS,用于挑选出一部分框。
1 | def filter_proposals(self, proposals, objectness, image_shapes, num_anchors_per_level): |
其中self._get_top_n_idx()
函数去取出概率最高的前 4000 个框的索引。最后的for
循环是根据特征图的框还原为原图的框,并选出最前面的 1000 个框(训练时是 2000 个,测试时是 1000 个)。
然后回到GeneralizedRCNN
的forward()
函数里的roi_heads()
,实际上是调用RoIHeads
,forward()
函数如下:
1 | def forward(self, features, proposals, image_shapes, targets=None): |
其中box_roi_pool()
是调用MultiScaleRoIAlign
把不同尺度的特征图池化到相同尺度,返回给box_features
,形状是\([1000, 256, 7, 7]\),1000 表示有 1000 个框(在训练时会从2000个选出 512 个,测试时则全部选,所以是 1000)。box_head()
是两个全连接层,返回的数形状是\([1000,1024]\),一个候选框使用一个 1024 的向量来表示。box_predictor()
输出最终的分类和边界框,class_logits
的形状是\([1000,91]\),box_regression
的形状是\([1000,364]\),\(364=91 \times 4\)。
然后回到GeneralizedRCNN
的forward()
函数中,transform.postprocess()
对输出进行后处理,将输出转换到原始图像的维度上。
下面总结一下 Faster RCNN 的主要组件:
- backbone
- rpn
- filter_proposals(NMS)
- rio_heads
下面的例子是使用 Faster RCNN 进行行人检测的 Finetune。数据集下载地址是https://www.cis.upenn.edu/~jshi/ped_html/,包括 70 张行人照片,345 个行人标签。
数据存放结构如下:
- PennFudanPed
- Annotation:标注文件,为
txt
- PedMasks:不清楚,没用到
- PNGImages:图片数据
- Annotation:标注文件,为
在Dataset
中,首先在构造函数中保存所有图片的文件名,后面用于查找对应的 txt 标签文件;在__getitem__()
函数中根据 index 获得图片和 txt 文件,查找 txt 文件的每一行是否有数字,有数字的则是带有标签的行,处理得到 boxes 和 labels,最后构造 target,target 是一个 dict,包括 boxes 和 labels。
在构造 DataLoader 时,还要传入一个collate_fn()
函数。这是因为在目标检测中,图片的宽高可能不一样,无法以 4D 张量的形式拼接一个 batch 的图片,因此这里使用 tuple 来拼接数据。
1 | # 收集batch data的函数 |
collate_fn 的输入是 list,每个元素是 tuple;每个 tuple 是 Dataset 中的 __getitem__()
返回的数据,包括(image, target)
举个例子:
1 | image=[1,2,3] |
输出为:
1 | batch: |
在代码中首先对数据和标签同时进行数据增强,因为对图片进行改变,框的位置也会变化,这里主要做了翻转图像和边界框的数据增强。
构建模型时,需要修改输出的类别为 2,一类是背景,一类是行人。
1 | model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True) |
这里不用构造 Loss,因为在 Faster RCNN 中已经构建了 Loss。在训练时,需要把 image 和 target 的 tuple 转换为 list,再输入模型。模型返回的不是真正的标签,而是直接返回 Loss,所以我们可以直接利用这个 Loss 进行反向传播。
代码如下:
1 | import os |
参考资料
如果你觉得这篇文章对你有帮助,不妨点个赞,让我有更多动力写出好文章。
我的文章会首发在公众号上,欢迎扫码关注我的公众号张贤同学。