Pytorch 学习笔记7

加载中...

本文可能存在错误，欢迎指正，共同学习进步。

有了前面的基础，我们就可以大致目标检测进行上手了，前面受限于篇幅，很多东西都没写完，寒假要是我记得一定埋坑！

示例参考TorchVision 目标检测微调教程进行操作，建议阅读原文。

这篇文章我们将使用Mask R-CNN模型来进行目标检测并通过微调来训练模型，我们将使用Penn-Fudan Database for Pedestrian Detection and Segmentation数据集，该数据集包含170个图像，345个行人实例，每个图像中的行人数目从2到10不等。

Mask R-CNN 概念

Mask R-CNN 可以参考原论文Mask R-CNN，也可以看看更通俗易懂的深度学习之目标检测的前世今生（Mask R-CNN）。

Mask R-CNN 是一种实例分割模型，它可以同时预测每个实例的边界框和掩码。Mask R-CNN 是一个两阶段的框架，第一阶段是 RPN，用于生成候选区域，第二阶段是 ROI 对齐，用于预测边界框和掩码。

Mask R-CNN是Faster R-CNN的扩展，它添加了一个分割头，允许我们预测每个实例的分割掩码。Mask R-CNN是一个实例分割模型，它可以同时预测每个实例的边界框和掩码。Mask R-CNN是一个两阶段的框架，第一阶段是RPN，用于生成候选区域，第二阶段是ROI对齐，用于预测边界框和掩码。

Mask R-CNN模型由三个主要部分组成：

backbone：用于提取特征的主干网络，通常是一个预训练的CNN模型。
RPN：区域提议网络，用于生成候选区域。
ROI 对齐：用于从候选区域中提取特征。

实验基础信息

系统环境：Windows 11, Pytorch latest, CUDA 12.4, python 3.11
硬件环境：RTX 3060Laptop, i7-12700H, 40GB RAM
模型：Mask R-CNN
数据集: Penn-Fudan 行人检测和分割数据库，包含170个图像，345个行人实例，每个图像中的行人数目从2到10不等。

实验基础运行环境

安装gautamchitnis的torchvision库

bash

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

bash

wget https://www.cis.upenn.edu/~jshi/ped_html/PennFudanPed.zip
unzip PennFudanPed.zip

bash

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

PedMasks目录包含了所有的实例分割掩码，PNGImages目录包含了所有的图像，Annotation目录包含了所有的标注文件，每个标注文件包含了对应图像中的行人实例的边界框和类别标签，我们随便看一个标注文件的内容：

txt

# Compatible with PASCAL Annotation Version 1.00
Image filename : "PennFudanPed/PNGImages/FudanPed00001.png"
Image size (X x Y x C) : 559 x 536 x 3
Database : "The Penn-Fudan-Pedestrian Database"
Objects with ground truth : 2 { "PASpersonWalking" "PASpersonWalking" }
# Note there may be some objects not included in the ground truth list for they are severe-occluded
# or have very small size.
# Top left pixel co-ordinates : (1, 1)
# Details for pedestrian 1 ("PASpersonWalking")
Original label for object 1 "PASpersonWalking" : "PennFudanPed"
Bounding box for object 1 "PASpersonWalking" (Xmin, Ymin) - (Xmax, Ymax) : (160, 182) - (302, 431)
Pixel mask for object 1 "PASpersonWalking" : "PennFudanPed/PedMasks/FudanPed00001_mask.png"

# Details for pedestrian 2 ("PASpersonWalking")
Original label for object 2 "PASpersonWalking" : "PennFudanPed"
Bounding box for object 2 "PASpersonWalking" (Xmin, Ymin) - (Xmax, Ymax) : (420, 171) - (535, 486)
Pixel mask for object 2 "PASpersonWalking" : "PennFudanPed/PedMasks/FudanPed00001_mask.png"

关于 labels 的一个说明。模型将类 0 视为背景。如果您的数据集不包含背景类，则您的 labels 中不应包含 0。例如，假设您只有两个类，猫和狗，您可以定义 1（而不是 0）表示猫，2 表示狗。因此，例如，如果其中一张图像同时包含这两个类，则您的 labels 张量应如下所示：[1, 2]。

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

实验步骤

数据集分析

在开始写代码之前，我们需要先分析一下数据集，我们需要知道数据集中包含了哪些信息，比如图像、标注文件、实例分割掩码等，我们需要知道数据集中的图像和标注文件是如何对应的，我们需要知道标注文件中包含了哪些信息，比如行人实例的边界框和类别标签等，我们需要知道实例分割掩码是如何生成的，我们需要知道实例分割掩码和标注文件是如何对应的等等。

比如，我读取了一个图像及其掩码图像，使用代码如下：

python

import matplotlib.pyplot as plt
from torchvision.io import read_image


image = read_image("data\PennFudanPed\PNGImages\FudanPed00016.png")
mask = read_image("data\PennFudanPed\PedMasks\FudanPed00016_mask.png")

plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)
plt.imshow(image.permute(1, 2, 0))
plt.axis("off")
plt.subplot(1, 2, 2)
plt.imshow(mask[0], cmap="gray")
plt.axis("off")
plt.show()

输出示例图

我们可以看到，左边是图像，右边是掩码图像，掩码图像是一个二值图像，其中白色部分表示行人实例，黑色部分表示背景。

然后我们再回过头看看标注文件，我们可以看到标注文件中包含了对应图像中的行人实例的边界框和类别标签，主要有以下几个重要的信息：

图像文件名
图像大小
行人实例的类别标签
行人实例的边界框
行人实例的掩码图像文件名

数据集类定义

接下来，我们需要定义一个数据集类，用于加载数据集，我们需要实现 __len__ 和 __getitem__ 方法，用于返回数据集的大小和数据集中的元素，我们还需要实现 get_height_and_width 方法，用于返回图像的高度和宽度。

python

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
        # 加载所有图像文件，并对它们进行排序以确保它们是排列整齐的
        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):
        # 加载图像和掩码
        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)
        # 实例被编码为不同的颜色
        obj_ids = torch.unique(mask)
        # 第一个ID是背景，所以移除它
        obj_ids = obj_ids[1:]
        num_objs = len(obj_ids)

        # 将彩色编码的掩码分割成一组二进制掩码
        masks = (mask == obj_ids[:, None, None]).to(dtype=torch.uint8)

        # 获取每个掩码的边界框坐标
        boxes = masks_to_boxes(masks)

        # 只有一个类别
        labels = torch.ones((num_objs,), dtype=torch.int64)

        image_id = idx
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        # 假设所有实例都不是密集型的（crowd）
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)

        # 将样本和目标包装到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)

以上的代码实现了一个数据集类，用于加载数据集，我们可以看到，我们首先加载图像和掩码，然后将掩码分割成一组二进制掩码，然后获取每个掩码的边界框坐标，然后将样本和目标包装到torchvision的tv_tensors中，最后返回图像和目标。接下来，我们试试上面的代码有没有成功读取到数据集，我们可以再添加一些代码来测试一下：

python


# 使用我们的数据集类
dataset = PennFudanDataset(root="./data/PennFudanPed", transforms=None)

import matplotlib.pyplot as plt
import matplotlib.patches as patches
import random

# 选择一个样本
idx = random.randint(0, len(dataset))
img, target = dataset[idx]

# 显示图像
plt.imshow(img.permute(1, 2, 0))
plt.axis("off")

# 显示掩码
for mask in target["masks"]:
    plt.imshow(mask, alpha=0.5, cmap="gray")

# 显示边界框
for box in target["boxes"]:
    x, y, w, h = box
    rect = patches.Rectangle(
        (x, y), w - x, h - y, linewidth=2, edgecolor="r", facecolor="none"
    )
    plt.gca().add_patch(rect)

plt.show()

可以看到输出结果如下所示：

输出示例图

Mask R-CNN 模型定义

Mask R-CNN

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

首先，假设我们从在COCO上预训练的Faster R-CNN模型开始，我们首先加载预训练的Faster R-CNN模型。

python

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

# 加载在COCO上预训练的Faster R-CNN模型 pretrained=True代表加载预训练模型
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights="DEFAULT")

# 获取分类器的输入特征数，一个分类+一个背景
num_classes = 2

# 获取分类器的输入特征数
in_features = model.roi_heads.box_predictor.cls_score.in_features

# 用新的头部替换预训练的头部
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

然后，我们对模型进行修改，使其支持不同的骨干网络：

python

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

# 加载一个预训练的模型，这里使用一个在COCO上预训练的模型
backbone = torchvision.models.mobilenet_v2(weights="DEFAULT").features

# FasterRCNN需要知道骨干网络的输出通道数量
# 对于mobilenet_v2，它是1280所以我们需要在这里添加
backbone.out_channels = 1280

# 定义AnchorGenerator
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),
                                   aspect_ratios=((0.5, 1.0, 2.0),))

# 定义ROI Pooling
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
                                                output_size=7,
                                                sampling_ratio=2)


# 将骨干网络、RPN、ROI Pooling、分类器和回归器组合到一个模型中
model = FasterRCNN(backbone,
                   num_classes=2,
                   rpn_anchor_generator=anchor_generator,
                   box_roi_pool=roi_pooler)

因为我们数据集仅含几百张图片，对于卷积神经网络来说，这数据集太小了，因此我们得站在巨人的肩膀上，使用预训练的模型来进行微调，这样可以加快训练速度，提高模型的准确性。

在这里，我们将使用一个在COCO数据集上预训练的模型，我们将加载预训练的模型，并将其修改为我们的数据集，然后对其进行微调。

python

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):
    # 加载一个在COCO上预训练的模型
    model = torchvision.models.detection.maskrcnn_resnet50_fpn(weights="DEFAULT")

    # 获取分类器的输入特征数
    in_features = model.roi_heads.box_predictor.cls_score.in_features

    # 用新的头部替换预训练的头部
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    # 现在获取掩码分类器的输入特征数
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = 256

    # 并用新的头部替换掩码预测器的头部
    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,hidden_layer,num_classes)

    return model

现在就实现他！

经过前面的各种准备，我们现在就可以实现目标检测了，我们需要定义一些辅助函数，比如计算准确率、计算损失函数等，然后我们就可以开始训练模型了。

首先，我们得定义一些辅助函数，不过这些辅助函数已经有前人实现过了，对于我们来说只需要下载调用就行了。

在windows上，我们首先要安装wget，再执行下载命令：

bash

winget install wget

然后我们就可以下载辅助函数了：

python

import os

os.system("wget https://raw.githubusercontent.com/pytorch/vision/main/references/detection/engine.py")
os.system("wget https://raw.githubusercontent.com/pytorch/vision/main/references/detection/utils.py")
os.system("wget https://raw.githubusercontent.com/pytorch/vision/main/references/detection/coco_utils.py")
os.system("wget https://raw.githubusercontent.com/pytorch/vision/main/references/detection/coco_eval.py")
os.system("wget https://raw.githubusercontent.com/pytorch/vision/main/references/detection/transforms.py")

上面这些还不够，我们再来定义一个数据增强函数：

python

from torchvision.transforms import v2 as T

# 数据增强函数
def get_transform(train):
    transforms = []
    if train:
        transforms.append(T.RandomHorizontalFlip(0.5))
    transforms.append(T.ToDtype(torch.float, scale=True))
    transforms.append(T.ToPureTensor())
    return T.Compose(transforms)

以上代码定义了一个数据增强函数，用于对图像和目标进行数据增强，我们可以看到，我们首先对图像进行随机水平翻转，然后将图像转换为float类型，最后将图像转换为tensor类型。

准备工作做完了，我们就可以开始训练模型了，我们首先加载数据集，然后定义数据增强函数，然后定义模型，然后定义优化器，然后定义损失函数，然后定义学习率调度器，然后定义训练器，然后开始训练模型。

我们先梳理一下基础流程：

代码如下：

python

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
        # 加载所有图像文件，并对它们进行排序以确保它们是排列整齐的
        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):
        # 加载图像和掩码
        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)
        # 实例被编码为不同的颜色
        obj_ids = torch.unique(mask)
        # 第一个ID是背景，所以移除它
        obj_ids = obj_ids[1:]
        num_objs = len(obj_ids)

        # 将彩色编码的掩码分割成一组二进制掩码
        masks = (mask == obj_ids[:, None, None]).to(dtype=torch.uint8)

        # 获取每个掩码的边界框坐标
        boxes = masks_to_boxes(masks)

        # 只有一个类别
        labels = torch.ones((num_objs,), dtype=torch.int64)

        image_id = idx
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        # 假设所有实例都不是密集型的（crowd）
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)

        # 将样本和目标包装到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)

from torchvision.transforms import v2 as T

# 定义训练和验证数据集，数据增强和标准化
def get_transform(train):
    transforms = []
    if train:
        transforms.append(T.RandomHorizontalFlip(0.5))
    transforms.append(T.ToDtype(torch.float, scale=True))
    transforms.append(T.ToPureTensor())
    return T.Compose(transforms)

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):
    # 加载在COCO上预训练的预训练模型
    model = torchvision.models.detection.maskrcnn_resnet50_fpn(weights="DEFAULT")

    # 获取分类器的输入特征数
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    # 用新的头部替换预先训练好的头部
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    # 现在获取掩码分类器的输入特征数
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = 256
    # 并用新的掩码预测器替换掩码预测器
    model.roi_heads.mask_predictor = MaskRCNNPredictor(
        in_features_mask, hidden_layer, num_classes
    )

    return model


from engine import train_one_epoch, evaluate
import utils


device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

# 我们的数据集只有两个类 - 行人和背景
num_classes = 2
# 使用我们的数据集和定义的转换
dataset = PennFudanDataset("data/PennFudanPed", get_transform(train=True))
dataset_test = PennFudanDataset("data/PennFudanPed", get_transform(train=False))

# 在训练和验证集中拆分数据集
indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices[:-50])
dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])

# 定义训练和验证数据加载器
data_loader = torch.utils.data.DataLoader(
    dataset, batch_size=2, shuffle=True, collate_fn=utils.collate_fn
)

data_loader_test = torch.utils.data.DataLoader(
    dataset_test, batch_size=1, shuffle=False, collate_fn=utils.collate_fn
)

# 获取模型实例，将其移动到设备，并构建优化器
model = get_model_instance_segmentation(num_classes)

model.to(device)

# 构建一个优化器.这里我们使用一个简单的SGD优化器，学习率：0.005，动量：0.9，权重衰减：0.0005
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)

# 学习率调度器
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

# 训练2个epochs
num_epochs = 10

for epoch in range(num_epochs):

    train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)
    # 更新学习率
    lr_scheduler.step()
    # 在测试集上评估
    evaluate(model, data_loader_test, device=device)

# 保存模型
torch.save(model.state_dict(), "model1.pth")

print("训练完成，模型已保存！")

以上的代码流程图如下所示：

在经过一段时间的训练后，我们就能得到 COCO 风格的 mAP > 50 - 65 的掩码 mAP，并且训练的模型已经保存在 model1.pth 文件中。

测试模型

我们可以使用以下代码来测试模型：

python

import torch
import torchvision
from PIL import Image
import torchvision.transforms as T
import matplotlib.pyplot as plt
import numpy as np
from pylab import mpl


mpl.rcParams["font.sans-serif"] = ["SimHei"]


# 加载模型
def get_model_instance_segmentation(num_classes):
    # 创建一个实例分割模型
    model = torchvision.models.detection.maskrcnn_resnet50_fpn(weights="DEFAULT")
    # 获取分类器的输入特征数
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    # 替换分类器
    model.roi_heads.box_predictor = (
        torchvision.models.detection.faster_rcnn.FastRCNNPredictor(
            in_features, num_classes
        )
    )
    # 获取掩码分类器的输入特征数
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = 256
    # 替换掩码预测器
    model.roi_heads.mask_predictor = (
        torchvision.models.detection.mask_rcnn.MaskRCNNPredictor(
            in_features_mask, hidden_layer, num_classes
        )
    )
    return model


# 定义类别数（背景+目标类别）
num_classes = 2

# 加载模型并加载训练好的权重
model = get_model_instance_segmentation(num_classes)
model.load_state_dict(torch.load("model1.pth"))
model.eval()

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model.to(device)

# 加载测试图像
img = Image.open("data/PennFudanPed/PNGImages/PennPed00093.png").convert("RGB")
# 图像预处理
transform = T.Compose(
    [
        T.ToTensor(),
    ]
)

img = transform(img)
img = img.to(device)

# 模型推理
with torch.no_grad():
    prediction = model([img])

# 可视化结果
img = img.cpu().permute(1, 2, 0).numpy()
img_np = img.copy()

plt.figure(figsize=(12, 8))

plt.subplot(1, 2, 1)
plt.imshow(img)
plt.axis("off")
plt.title("原图")

plt.subplot(1, 2, 2)
plt.imshow(img)

# 获取预测结果
boxes = prediction[0]["boxes"].cpu().numpy()
scores = prediction[0]["scores"].cpu().numpy()
masks = prediction[0]["masks"].cpu().numpy()

# 只保留得分高于阈值的预测结果
threshold = 0.5
indices = scores > threshold
boxes = boxes[indices]
masks = masks[indices]

# 在右侧子图上叠加掩码和边界框
for box, mask in zip(boxes, masks):
    x1, y1, x2, y2 = box

    # 叠加掩码（使用红色半透明效果）
    mask = mask[0]
    masked_area = np.zeros_like(img)
    masked_area[:, :, 0] = mask * 1.0
    plt.imshow(masked_area, alpha=0.5, cmap="Reds")

    # 绘制边界框
    plt.gca().add_patch(
        plt.Rectangle(
            (x1, y1), x2 - x1, y2 - y1, linewidth=2, edgecolor="red", facecolor="none"
        )
    )

plt.axis("off")
plt.title("预测结果")
plt.tight_layout()
plt.show()

可以看到，识别结果如下所示：

识别结果

Pytorch 学习笔记7https://blog.tianli0.top/posts/2024/Pytorch7

作者Tianli

发布于12/11

更新于今日内

许可协议 CC BY-NC-SA 4.0

署名-非商业性使用-相同方式共享 4.0 国际

深度学习学习笔记 Pytorch CUDA

反馈与投诉

评论隐私政策

Mask R-CNN 概念 ​

实验基础信息 ​

实验基础运行环境 ​

实验步骤 ​

数据集分析 ​

数据集类定义 ​

Mask R-CNN 模型定义 ​

现在就实现他！ ​

测试模型 ​

Mask R-CNN 概念

实验基础信息

实验基础运行环境

实验步骤

数据集分析

数据集类定义

Mask R-CNN 模型定义

现在就实现他！

测试模型