pytorch深度学习入门（9）之-使用TORCHVISION进行目标检测

原创

已于 2023-11-30 11:09:36 修改 · 3.1k 阅读

6 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #pytorch #目标检测

于 2023-11-14 11:49:36 首次发布

本教程介绍了如何利用torchvision库进行目标检测，特别是使用预训练的Mask R-CNN模型在PennFudan数据集上进行微调。内容涵盖了数据集的定义、模型的构建和微调，展示了在自定义数据集上训练和评估对象检测模型的过程。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

概述：

torchvision是PyTorch项目的一部分，它是一个用于构建计算机视觉模型的图形库。torchvision提供了以下功能：

数据集：torchvision提供了许多常用的数据集，这些数据集适用于训练和测试计算机视觉模型。
模型架构：torchvision包含许多常用的计算机视觉模型架构，例如ResNet、VGG、AlexNet等。
图像转换：torchvision还提供了许多用于图像转换的函数，例如裁剪、旋转等。
其他工具：torchvision还包含一些其他有用的工具，例如用于模型训练和测试的数据加载器等。
总之，torchvision是一个功能强大的计算机视觉库，适用于构建各种类型的计算机视觉模型，包括图像分类、目标检测、图像分割等。

在本教程中，我们将在Penn-Fudan 行人检测和分割数据库上微调预训练的Mask R-CNN模型。它包含 170 张图像和 345 个行人实例，我们将用它来说明如何使用 torchvision 中的新功能，以便在自定义数据集上训练对象检测和实例分割模型。

定义数据集

用于训练对象检测、实例分割和人物关键点检测的参考脚本可以轻松支持添加新的自定义数据集。数据集应继承自标准 torch.utils.data.Dataset类，并实现__len__和 getitem。

我们需要的唯一特殊性是数据集__getitem__ 应返回一个元组：

image：torchvision.tv_tensors.Imageshape 、纯张量或 size 的 PIL 图像[3, H, W](H, W)

target：包含以下字段的字典

boxes，torchvision.tv_tensors.BoundingBoxes形状：[N, 4]
N边界框的坐标格式，范围为 to和to[x0, y0, x1, y1]0W0H

labels，torch.Tensor形状整数[N]：每个边界框的标签。
0始终代表背景类。

image_id, int：图像标识符。它应该是
在数据集中的所有图像之间是唯一的，并且在评估期间使用

area，torch.Tensor形状的浮动[N]：边界框的面积。这是用的
在使用 COCO 度量进行评估时，将小、中、大盒子之间的度量分数分开。

iscrowd, uint8torch.Tensor形状[N]: 带有iscrowd=Truewill 的实例
评估期间被忽略。

（可选）masks，torchvision.tv_tensors.Mask形状：分割[N, H, W]
每个对象的掩模

如果您的数据集符合上述要求，那么它将适用于参考脚本中的训练和评估代码。评估代码将使用 pycocotools可以通过.

pip install pycocotools

进行安装
注意：

=== 对于 Windows，请使用命令pycocotools从gautamchitnis安装 ===

pip install git+https://github.com/gautamchitnis/cocoapi.git@cocodataset-master#subdirectory=PythonAPI

关于的一个注释labels。该模型将阶级视为0背景。如果您的数据集不包含背景类，则您0的labels. 例如，假设您只有两个类：cat和dogs，您可以定义1(not 0) 来表示cats并2表示dogs。因此，例如，如果其中一张图像同时具有这两个类，则您的labels张量应该类似于。[1, 2]

此外，如果您想在训练期间使用宽高比分组（以便每个批次仅包含具有相似宽高比的图像），那么建议还实现一个方法，该方法返回图像的高度和宽度get_height_and_width 。如果未提供此方法，我们将通过查询数据集的所有元素 getitem，这会将图像加载到内存中，并且比提供自定义方法要慢。

为 PennFudan 编写自定义数据集
让我们为 PennFudan 数据集编写一个数据集。下载并解压 zip 文件后，我们有以下文件夹结构：

PennFudanPed/
    PedMasks/
        FudanPed00001_mask.png
        FudanPed00002_mask.png
        FudanPed00003_mask.png
        FudanPed00004_mask.png
        ...
    PNGImages/
        FudanPed00001.png
        FudanPed00002.png
        FudanPed00003.png
        FudanPed00004.png

这是一对图像和分割掩模的一个示例

../_static/img/tv_tutorial/tv_image01.png ../_static/img/tv_tutorial/tv_image02.png
在这里插入图片描述

因此，每个图像都有一个相应的分割掩模，其中每种颜色对应于不同的实例。让我们torch.utils.data.Dataset为这个数据集编写一个类。在下面的代码中，我们将图像、边界框和掩模包装到 torchvision.TVTensor类中，以便我们能够为给定的对象检测和分割任务应用 torchvision 内置转换（新的 Transforms API ）。也就是说，图像张量将被包裹torchvision.tv_tensors.Image，将边界框包裹为 torchvision.tv_tensors.BoundingBoxes，将掩模包裹为torchvision.tv_tensors.Mask。torchvision.TVTensor与子类一样torch.Tensor，包装对象也是张量并继承普通 torch.TensorAPI。有关 torchvision 的更多信息，tv_tensors请参阅此文档。

import os
import torch

from torchvision.io import read_image
from torchvision.ops.boxes import masks_to_boxes
from torchvision import tv_tensors
from torchvision.transforms.v2 import functional as F


class PennFudanDataset(torch.utils.data.Dataset):
    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 = read_image(img_path)
        mask = read_image(mask_path)
        # instances are encoded as different colors
        obj_ids = torch.unique(mask)
        # first id is the background, so remove it
        obj_ids = obj_ids[1:]
        num_objs = len(obj_ids)

        # split the color-encoded mask into a set
        # of binary masks
        masks = (mask == obj_ids[:, None, None]).to(dtype=torch.uint8)

        # get bounding box coordinates for each mask
        boxes = masks_to_boxes(masks)

        # there is only one class
        labels = torch.ones((num_objs,), dtype=torch.int64)

        image_id = idx
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        # suppose all instances are not crowd
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)

        # Wrap sample and targets into torchvision tv_tensors:
        img = tv_tensors.Image(img)

        target = {
   
   }
        target["boxes"] = tv_tensors.BoundingBoxes(boxes, format="XYXY", canvas_size=F.get_size(img))
        target["masks"] = tv_tensors.Mask(masks)
        target["labels"] = labels
        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)

这就是数据集的全部内容。现在让我们定义一个可以对此数据集执行预测的模型。

定义你的模型

在本教程中，我们将使用Mask R-CNN，它基于 Faster R-CNN。Faster R-CNN 是一种预测图像中潜在对象的边界框和类别分数的模型。

../_static/img/tv_tutorial/tv_image03.png

Mask R-CNN 在 Faster R-CNN 中添加了一个额外的分支，它还预测每个实例的分割掩模。

../_static/img/tv_tutorial/tv_image04.png

在两种常见情况下，人们可能想要修改 TorchVision Model Zoo 中的可用模型之一。第一种是当我们想要从预先训练的模型开始，并对最后一层进行微调时。另一种是当我们想要用不同的模型替换模型的主干时（例如，为了更快的预测）。

让我们看看在接下来的部分中我们将如何做其中一个或另一个。

1 - 从预训练模型进行微调

假设您想从在 COCO 上预训练的模型开始，并希望针对您的特定类别对其进行微调。这是一种可能的方法：

import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor

# load a model pre-trained on COCO
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights="DEFAULT")

# replace the classifier with a new one, that has
# num_classes which is user-defined
num_classes = 2  # 1 class (person) + background
# get number of input features for the classifier
in_features = model.roi_heads.box_predictor.cls_score.in_features
# replace the pre-trained head with a new one
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

2 - 修改模型以添加不同的主干

import torchvision
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator

# load a pre-trained model for classification and return
# only the features
backbone = torchvision.models.mobilenet_v2(weights="DEFAULT").features
# ``FasterRCNN`` needs to know the number of
# output channels in a backbone. For mobilenet_v2, it's 1280
# so we need to add it here
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
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 = torchvision.ops.MultiScaleRoIAlign(
    featmap_names=['0'],
    output_size=7,
    sampling_ratio=2,
)

# put the pieces together inside a Faster-RCNN model
model = FasterRCNN(
    backbone,
    num_classes=2,
    rpn_anchor_generator=anchor_generator,
    box_roi_pool=roi_pooler,
)

PennFudan 数据集的对象检测和实例分割模型

在我们的例子中，我们希望从预训练的模型中进行微调，因为我们的数据集非常小，所以我们将遵循方法 1。

在这里，我们还想计算实例分割掩码，因此我们将使用 Mask R-CNN：

import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor


def get_model_instance_segmentation(num_classes):
    # load an instance segmentation model pre-trained on COCO
    model = torchvision.models.detection.maskrcnn_resnet50_fpn(weights="DEFAULT")

    # get number of input features for the classifier
    in_features