本文为学习笔记
英文教程地址：TorchVision Object Detection Finetuning Tutorial — PyTorch Tutorials 1.8.1+cu102 documentation
中文翻译地址：PyTorch NLP From Scratch: 使用char-RNN对姓氏进行分类_w3cschool
# 引入库
import os
import numpy as np
import torch
from PIL import Image
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
from engine import train_one_epoch, evaluate
import utils
import transforms as Tclass PennFudanDataset(object):
def __init__(self, root, transforms):
self.root = root
self.transforms = transforms
# load all image files, sorting them to
# ensure that they are aligned
self.imgs = list(sorted(os.listdir(os.path.join(root, "PNGImages"))))
self.masks = list(sorted(os.listdir(os.path.join(root, "PedMasks"))))
def __getitem__(self, idx):
# load images and masks
img_path = os.path.join(self.root, "PNGImages", self.imgs[idx])
mask_path = os.path.join(self.root, "PedMasks", self.masks[idx])
img = Image.open(img_path).convert("RGB")
# note that we haven't converted the mask to RGB,
# because each color corresponds to a different instance
# with 0 being background
mask = Image.open(mask_path)
# convert the PIL Image into a numpy array
mask = np.array(mask)
# instances are encoded as different colors
obj_ids = np.unique(mask)
# first id is the background, so remove it
obj_ids = obj_ids[1:]
# split the color-encoded mask into a set
# of binary masks
masks = mask == obj_ids[:, None, None]
# get bounding box coordinates for each mask
num_objs = len(obj_ids)
boxes = []
for i in range(num_objs):
pos = np.where(masks[i])
xmin = np.min(pos[1])
xmax = np.max(pos[1])
ymin = np.min(pos[0])
ymax = np.max(pos[0])
boxes.append([xmin, ymin, xmax, ymax])
# convert everything into a torch.Tensor
boxes = torch.as_tensor(boxes, dtype=torch.float32)
# there is only one class
labels = torch.ones((num_objs,), dtype=torch.int64)
masks = torch.as_tensor(masks, dtype=torch.uint8)
image_id = torch.tensor([idx])
area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
# suppose all instances are not crowd
iscrowd = torch.zeros((num_objs,), dtype=torch.int64)
target = {}
target["boxes"] = boxes
target["labels"] = labels
target["masks"] = masks
target["image_id"] = image_id
target["area"] = area
target["iscrowd"] = iscrowd
if self.transforms is not None:
img, target = self.transforms(img, target)
return img, target
def __len__(self):
return len(self.imgs)### 简单介绍
# 自己搭建神经网络时，一般都采用已有的网络模型，在其基础上进行修改
# 在两种常见情况下，可能要修改 Torchvision modelzoo 中的模型。
# 第一种是当我们想使用预训练模型，并微调最后一层 (roi_heads) 时
# 另一种是当我们要用另一个模型替换预定义模型的主干 (backbone) 时(例如，为了更快的预测）。
## 第一种：微调已经预训练过的模型
# 使用的模型是在 COCO 上预先训练过的模型 fasterrcnn_resnet50_fpn
# 该模型是一个 FasterRCNN 模型，其网络结构为：transform、backbone、rpn、roi_heads
# FasterRCNN 是一个目标检测网络，推荐先了解 FasterRCNN 后再看后面的代码
# 参考文章
# transform: GeneralizedRCNNTransform 层对输入的图像进行大小变换，用于标准化训练集，方便模型训练
# backbone: BackboneWithFPN 层用于对图像进行特征提取，特征提取得到 feature_map，后续操作通过 feature_map 上操作
# rpn(region proposal network): FeaturePyramidNetwork 层用于获取建议框，建议框是模型粗略框取的图片中的物品区域，建议框的大小并不固定
# roi_heads: RoIHeads 层包括 box_roi_pool、box_head、box_predictor 三层，用于最后处理得到预测结果
# -- box_roi_pool 将由建议框截取的大小不一的 feature_map 信息进行标准化，得到同等大小的局部特征层
# -- box_head
# -- box_predictor 根据局部特征层进行分类预测和回归预测包括分类器 cls_score 和回归器 bbox_pred
# 分类器 cls_score 用于区分每个框取的对象是什么、回归器 bbox_pred 用于确定对每个对象框取的预测框的框坐标和框的长宽共4个量的值
# torchvision 是 pytorch 中用于处理图像视频的工具集合
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
# 用具有用户定义的 num_classes 的新分类器替换原分类器
# 这里指的分类器是 FasterRCNN 模型最后一层结构 roi_heads 的第三个子层 box_predictor，也可以叫预测器
# 本来模型的分类类别数（num_classes）为 91，现在将其改为 2
# bbox_pred 的 out_feature 为 cls_score 的 4 倍，这是因为 bbox_pred 输出的是预测框的数据
# 预测框包含一个坐标数据（x,y）以及长宽（width，height），共 4 个量
# 可见每个 cls_score 的输出对应 4 个 bbox_pred 的输出
'''
原分类器为：
FastRCNNPredictor(
(cls_score): Linear(in_features=1024, out_features=91, bias=True)
(bbox_pred): Linear(in_features=1024, out_features=364, bias=True)
)
替换后的分类器为：
FastRCNNPredictor(
(cls_score): Linear(in_features=1024, out_features=2, bias=True)
(bbox_pred): Linear(in_features=1024, out_features=8, bias=True)
)
'''
num_classes = 2
# 获取分类器的输入要素数量 in_features = 1024，以保证替换后的分类器能够接收上层的数据
# mode.roi_heads.box_predictor 是 fasterrcnn_resnet50_fpn 的最后一层结构 roi_heads 的子结构 box_predictor
# box_predictor 包含两层：cls_score、bbox_pred，这里需要替换的是 cls_score
# (cls_score): Linear(in_features=1024, out_features=91, bias=True)
in_features = model.roi_heads.box_predictor.cls_score.in_features
# 用新的 FastRCNNPredictor 分类器替换原分类器
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
## 第二种：修改预定义模型的主干
# 修改主干即修改模型的 backbone 部分
import torchvision
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator
# 加载预先训练的模型进行分类，仅返回 features，(模型 mobilenet_v2 的 backbone 定义在 features 中)
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
# backbone.out_channels 记录 backbone 层输出的数据维度，对于 mobilenet_v2 来说，输出维度为 1280
backbone.out_channels = 1280
# let's make the RPN generate 5 x 3 anchors per spatial location, with 5 different sizes and 3 different aspect ratios.
# We have a Tuple[Tuple[int]] because each feature map could potentially have different sizes and aspect ratios
'''
rpn 层生成建议框的过程：
首先，生成锚点框，可以理解为在每个点上生成几个不同大小的框，由此一张图上会有很多很多锚点框（anchor）
然后，筛选锚点框，留下框住了物体的框以及调整这些框大小，得到建议框
建议框交给下一级进一步处理
'''
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),
aspect_ratios=((0.5, 1.0, 2.0),))
# let's define what are the feature maps that we will
# use to perform the region of interest cropping, as well as
# the size of the crop after rescaling.
# if your backbone returns a Tensor, featmap_names is expected to
# be [0]. More generally, the backbone should return an
# OrderedDict[Tensor], and in featmap_names you can choose which
# feature maps to use.
# roi_pooler 是一个 box_roi_pool，将由建议框截取的大小不一的 feature_map 信息进行标准化，得到同等大小的局部特征层
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
# 将以上定义的零散结构组合成一个 FasterRCNN
model = FasterRCNN(backbone,
num_classes=2,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
本文章为转载内容，我们尊重原作者对文章享有的著作权。如有内容错误或侵权问题，欢迎原作者联系我们进行内容更正或删除文章。
