1
2
3
4

obj_ids = torch.unique(mask)  
# first id is the background, so remove it  
obj_ids = obj_ids[1:]  
num_objs = len(obj_ids)

1
2
3
4

mask = tensor([[ 0,  0, 10, 10,  0],
               [ 0, 10, 10, 10,  0],
               [ 0, 20, 20,  0,  0],
               [20, 20, 20,  0,  0]])

1

masks = (mask == obj_ids[:, None, None]).to(dtype=torch.uint8)

• 功能: 这行代码的作用是将一个多类别、颜色编码的掩码图像（mask）转换成一组二进制掩码
• 解释:
  • mask：通常是一个二维数组（代表图像），其中每个像素的值代表一个物体的ID或者类别。背景通常为0。
  • obj_ids：这是一个包含图像中所有物体唯一ID的一维张量。
  • obj_ids[:, None, None]：这是一种广播（broadcasting）技巧。它将一维的 obj_ids 扩展成三维，以便能和二维的 mask 进行逐像素比较。
  • (mask == obj_ids[:, None, None])：这个比较操作会为 obj_ids 中的每一个物体ID生成一个布尔类型的二进制掩码。如果 mask 中的像素值等于某个物体ID，那么在对应的二进制掩码中，该像素位置为 True，否则为 False。最终生成的 masks 是一个三维张量，形状为 [N, H, W]，其中 N 是物体的数量，H 和 W 是图像的高和宽。
  • .to(dtype=torch.uint8)：将布尔类型的掩码转换为无符号8位整型（0和1），这是PyTorch中常用的掩码格式。

1
2

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

• 功能: 这行代码调用 masks_to_boxes 函数，为上一部生成的每个二进制掩码计算出其最小外接边界框。
• 解释:
  • masks_to_boxes 是 torchvision.ops 模块中的一个函数。
  • 它接收一个形状为 [N, H, W] 的掩码张量，并返回一个形状为 [N, 4] 的张量，其中每一行代表一个边界框，格式为 (x1, y1, x2, y2)，即左上角和右下角的坐标。

1
2

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

• 功能: 为每个检测到的物体分配一个标签。
• 解释:
  • num_objs：图像中物体的数量。
  • torch.ones((num_objs,), dtype=torch.int64)：创建一个长度为 num_objs 的张量，所有元素都为1。这行代码假设数据集中只有一个物体类别（不包括背景）。在实际应用中，如果存在多个类别，你需要根据具体情况为每个物体分配正确的类别标签。

1
2
3
4

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

• 功能: 计算一些与COCO数据集格式兼容的元数据。
• 解释:
  • image_id: 为当前图像分配一个唯一的ID，这里直接使用了数据加载时的索引 idx。
  • area: 计算每个边界框的面积。这在一些评估指标中会用到。
  • iscrowd: 一个标志，用于指示某个物体是否为“拥挤”的（即由多个物体组成的区域）。这里假设所有物体都不是拥挤的，所以都设置为0。

1
2
3
4
5
6
7
8
9
10

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

• 功能: 将图像和所有标注信息打包成torchvision的tv_tensors格式，并存入一个字典target中。
• 解释:
  • tv_tensors 是 torchvision.transforms.v2 引入的新特性，它们是 torch.Tensor 的子类，能够携带元数据（如边界框的格式、图像尺寸等），并在进行数据增强时自动更新这些元数据。
  • tv_tensors.Image(img): 将普通的图像张量包装成Image类型的tv_tensor。
  • tv_tensors.BoundingBoxes(...): 将边界框张量包装成BoundingBoxes类型的tv_tensor。 这需要提供边界框的格式（这里是XYXY）和原始图像的尺寸（canvas_size），以便在图像变换时能正确地调整边界框。
  • tv_tensors.Mask(masks): 将掩码张量包装成Mask类型的tv_tensor。
  • 最后，将所有这些tv_tensors和其他元数据存入名为target的字典中。这个字典是许多torchvision模型的标准输入格式。

1
2
3
4

if self.transforms is not None:
    img, target = self.transforms(img, target)

return img, target

• 功能: 如果定义了数据增强操作（self.transforms），则将其应用于图像和相应的target。
• 解释:
  • self.transforms 通常是一个包含一系列数据增强操作（如随机裁剪、翻转、缩放等）的组合。
  • 由于图像和标注（边界框、掩码）都被包装成了tv_tensors，当对img进行几何变换时，target中的boxes和masks也会被自动地、正确地进行相应的变换。
  • 最后，返回经过（可能有的）数据增强处理后的图像和标注。

第1步：加载预训练模型

1
2
3
4
5

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

# 加载一个在COCO上预训练过的模型
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights="DEFAULT")

• import ...: 导入所需的库和模块。torchvision 是PyTorch官方的计算机视觉库，FastRCNNPredictor 是一个辅助类，用于方便地替换Faster R-CNN模型的预测头。
• torchvision.models.detection.fasterrcnn_resnet50_fpn(...): 这行代码加载了一个非常经典和强大的目标检测模型——Faster R-CNN。
  • resnet50: 表示这个模型使用ResNet-50作为其“骨干网络（backbone）”。骨干网络负责从输入图像中提取有用的特征。ResNet-50是一个非常强大的图像特征提取器。
  • fpn: 表示模型使用了特征金字塔网络（Feature Pyramid Network）。FPN能够整合骨干网络在不同层次提取的特征，从而有效地检测不同尺寸的物体（大的和小的都能很好地识别）。
• weights="DEFAULT": 这是最关键的部分。它告诉torchvision加载在COCO数据集上已经训练好的权重。COCO是一个包含80个物体类别的大型、多样化的数据集。这意味着我们得到的model已经是一个“专家”了，它已经学会了如何识别各种通用的视觉特征，比如边缘、纹理、形状，甚至是一些常见的物体部件。

第2步：为新任务定义类别数

1
2

# 将分类器替换为一个新的、具有用户自定义类别数的分类器
num_classes = 2  # 1 个类别 (例如 '人') + 背景

• num_classes = 2: 这里我们定义了我们自己任务的类别数量。
• 为什么是2，而不是1？: 在目标检测任务中，模型不仅需要识别出你感兴趣的物体，还需要能识别出哪些区域是“背景”。因此，总的类别数总是 你感兴趣的物体类别数 + 1 (背景)。如果你要检测猫和狗，那么 num_classes 就应该是 2 + 1 = 3。

第3步：获取并替换模型的“头部”

1
2
3
4

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

• model.roi_heads.box_predictor.cls_score: 这是深入到模型内部结构，找到负责最终类别预测的那个线性层（Linear Layer）。
  • roi_heads：是Faster R-CNN中负责处理候选区域（Region of Interest）的“头”部网络。
  • box_predictor：是头部网络中具体执行分类和边界框回归预测的模块。
  • cls_score：是box_predictor中最终输出类别分数的那个全连接层。
• .in_features: 这个属性可以获取到cls_score这个线性层所期望的输入向量的维度（大小）。我们必须知道这个数字，才能创建一个可以无缝衔接上去的新层。
• model.roi_heads.box_predictor = FastRCNNPredictor(...): 这是整个替换操作的核心。
  • 我们创建了一个全新的FastRCNNPredictor实例。
  • FastRCNNPredictor(in_features, num_classes)：这个新的预测器被初始化为：
    1. 接收与原始模型完全相同大小的输入特征（in_features）。
    2. 但它的输出维度是我们新任务所需的类别数（num_classes，也就是2）。
  • 通过这个赋值操作，我们把模型原来那个为COCO数据集（比如91个类别）设计的“旧头”给扔掉了，换上了我们自己为新任务（2个类别）设计的“新头”。

这样做的好处是，我们不需要从零开始训练整个庞大的网络，只需要在我们自己的（通常较小的）数据集上，主要训练这个新的、小得多的“头部”，以及微调一下模型的其他部分。这极大地加快了训练速度，降低了对数据量的要求，并且通常能达到比从零开始训练好得多的效果。

前一个部分是拿一个预制好的完整模型然后只换掉“头”，而这个部分则是我们自己挑选“骨干”、“锚点生成器”和“感兴趣区域池化层”这些核心组件，然后把它们拼装起来。

这个方法提供了极大的灵活性，你可以任意替换模型的某个部分（比如用一个更轻量级的骨干网络），而不需要重写整个模型。

第1步：选择并准备骨干网络 (Backbone)

1
2
3
4
5

# 加载一个预训练的分类模型，并且只返回其特征提取部分
backbone = torchvision.models.mobilenet_v2(weights="DEFAULT").features
# FasterRCNN 需要知道骨干网络的输出通道数。对于mobilenet_v2，这个值是1280
# 所以我们需要在这里手动添加它
backbone.out_channels = 1280

• 功能: 定义模型的“眼睛”，即特征提取器。
• 解释:
  • torchvision.models.mobilenet_v2(weights="DEFAULT"): 这里我们选择 MobileNetV2 作为骨干网络，而不是像之前例子里的 ResNet-50。MobileNetV2 是一个非常轻量级且高效的网络，非常适合在移动设备或资源有限的环境下运行。我们同样加载了在 ImageNet 上预训练好的权重。
  • .features: 这是关键。完整的 MobileNetV2 模型最后有一个用于分类的“头”（一个全连接层）。通过 .features，我们把这个头给“砍掉”了，只保留了前面的卷积层部分。因为在目标检测中，我们只需要骨干网络来提取特征图（Feature Map），而不需要它来进行最终的图像分类。
  • backbone.out_channels = 1280: 这是一个非常重要的手动设置。FasterRCNN 的构造函数需要知道从骨干网络输出的特征图的“深度”（即通道数）是多少，以便正确地连接后续的RPN网络。对于标准的 MobileNetV2，其 .features 部分的输出通道数是1280。因为我们得到的是一个通用的 nn.Sequential 模块，它本身没有 .out_channels 这个属性，所以我们必须手动为它指定。

第2步：自定义锚点生成器 (Anchor Generator)

1
2
3
4
5
6

# 让RPN网络在特征图的每个空间位置上生成 5x3=15 个锚点
# 这15个锚点由5种不同的尺寸和3种不同的宽高比组合而成
anchor_generator = AnchorGenerator(
    sizes=((32, 64, 128, 256, 512),),
    aspect_ratios=((0.5, 1.0, 2.0),)
)

• 功能: 定义区域提议网络（RPN）用来产生候选框的“模板”。
• 解释:
  • RPN 不会盲目地在图像上搜索物体。它会在特征图的每一个位置上，放置一堆预设好的、不同尺寸和形状的“锚点框”（Anchors）。然后，网络会学习判断这些锚点框里是否包含物体，并对框的位置进行微调。
  • sizes=((32, 64, 128, 256, 512),): 定义了锚点框的5种基本尺寸（以像素为单位）。注意这里的双括号 ((...))，这是因为模型可以处理来自特征金字塔（FPN）的多个特征图，每个特征图可以有自己的一套尺寸。由于我们的 MobileNetV2 backbone只输出一个特征图，所以我们只提供了一组尺寸。
  • aspect_ratios=((0.5, 1.0, 2.0),): 定义了3种宽高比。1.0 是正方形，0.5 是矮胖的矩形，2.0 是瘦高的矩形。
  • 最终，在特征图的每个点，都会生成 5 (尺寸) * 3 (宽高比) = 15 个锚点框。通过自定义这些值，你可以让模型更好地适应你数据集中物体的特定形状和大小。

第3步：定义感兴趣区域池化层 (ROI Pooler)

1
2
3
4
5
6

# 定义我们将从哪些特征图上进行感兴趣区域(RoI)裁剪，以及裁剪后缩放的尺寸。
roi_pooler = torchvision.ops.MultiScaleRoIAlign(
    featmap_names=['0'],
    output_size=7,
    sampling_ratio=2
)

• 功能: 从特征图中精确地提取每个候选区域的特征，并将其缩放到一个固定的尺寸。
• 解释:
  • 在RPN提出了很多大小不一的候选区域（Region of Interest, RoI）后，我们需要从这些区域中提取特征，送入最终的分类和边界框回归“头”。但这个“头”要求输入尺寸是固定的。RoIAlign 就是来做这个“裁剪并缩放”的工作的。
  • featmap_names=['0']: 这个参数至关重要，它告诉 RoIAlign 该从哪一张特征图上提取特征。因为我们的 backbone 只输出一个张量（Tensor），而不是一个带有名字的有序字典（OrderedDict），torchvision 会默认将其命名为 '0'。如果我们用的是FPN，这里可能会是 ['0', '1', '2', '3'] 这样的列表。
  • output_size=7: 无论原始的候选区域有多大，从其中提取的特征图都将被缩放到 7x7 的大小。
  • sampling_ratio=2: 这是 RoIAlign 算法的一个参数，用于更精确地对齐特征，避免量化误差，从而提升检测精度。

第4步：组装模型

1
2
3
4
5
6
7

# 将所有部件组装进一个FasterRCNN模型里
model = FasterRCNN(
    backbone,
    num_classes=2,
    rpn_anchor_generator=anchor_generator,
    box_roi_pool=roi_pooler
)

• 功能: 调用 FasterRCNN 的构造函数，将我们上面精心准备的所有自定义组件传入，完成模型的创建。
• 解释:
  • backbone: 我们选择的 MobileNetV2 特征提取器。
  • num_classes=2: 和上一个例子一样，代表1个物体类别 + 1个背景类别。
  • rpn_anchor_generator=anchor_generator: 传入我们自定义的锚点生成器。
  • box_roi_pool=roi_pooler: 传入我们自定义的ROI池化层。

这个是针对**实例分割（Instance Segmentation）**任务，使用的是 Mask R-CNN 模型。

实例分割比目标检测更进了一步，它不仅要找出物体的位置（边界框），还要为每个物体实例生成一个像素级的掩码（mask）。这里的逻辑和微调 Faster R-CNN 非常相似，但由于 Mask R-CNN 多了一个“生成掩码”的分支，所以我们需要替换两个“头”：一个是负责分类的头，另一个是负责生成掩码的头。

第1步：加载预训练的 Mask R-CNN 模型

1
2
3

def get_model_instance_segmentation(num_classes):
    # 加载一个在COCO上预训练过的实例分割模型
    model = torchvision.models.detection.maskrcnn_resnet50_fpn(weights="DEFAULT")

• 功能: 加载一个强大的、预训练好的实例分割模型。
• 解释:
  • torchvision.models.detection.maskrcnn_resnet50_fpn: 这次我们加载的是 Mask R-CNN 模型。它的结构可以看作是在 Faster R-CNN 的基础上增加了一个并行的分支，这个分支专门用来为每个检测到的物体（RoI）生成二进制掩码。
  • resnet50_fpn: 同样，它使用 ResNet-50 作为骨干网络，并带有 FPN，这对于在不同尺度上同时做好目标检测和实例分割至关重要。
  • weights="DEFAULT": 加载在 COCO 数据集上预训练的权重。这个预训练模型不仅知道如何检测80个类别的物体，还知道如何为它们生成精确的像素级掩码。

第2步：替换边界框分类头 (Box Predictor)

1
2
3
4

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

• 功能: 调整模型以适应我们新任务的物体类别。
• 解释:
  • 这部分和微调 Faster R-CNN 的代码完全一样。
  • model.roi_heads 是处理候选区域（RoI）的总模块。
  • 它内部的 box_predictor 负责两件事：1) 判定候选区域属于哪个类别；2) 微调候选区域的边界框位置。
  • 我们通过 cls_score 找到负责分类的那个全连接层，获取其输入维度 in_features。
  • 然后，我们创建一个新的 FastRCNNPredictor，它的输出维度是我们自定义的 num_classes（例如，你的物体数 + 1个背景）。
  • 这个操作将原来为COCO数据集（91个类别）设计的分类器，换成了我们自己任务专用的分类器。

第3步：替换掩码预测头 (Mask Predictor)

1
2
3
4
5
6
7
8
9

# 现在获取掩码分类器的输入特征数
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
)

• 功能: 调整模型以生成与我们新任务类别相对应的掩码。
• 解释: 这是 Mask R-CNN 特有的步骤。
  • model.roi_heads.mask_predictor: 这是 roi_heads 中与 box_predictor 并行存在的另一个模块，它专门负责生成掩码。
  • conv5_mask: 我们深入到 mask_predictor 内部，找到它的一个关键卷积层（在这里是conv5_mask）来获取其输入通道数 in_channels。这个值决定了输入到掩码头的特征图的深度。
  • hidden_layer = 256: Mask R-CNN的掩码头通常由一系列卷积层组成。这里我们定义了这些中间卷积层的通道数（或称特征维度）为256。这是一个常见的默认值。
  • MaskRCNNPredictor(...): 我们创建了一个新的掩码预测器。
    • in_features_mask: 它的输入通道数必须和原始模型匹配，确保可以无缝衔接。
    • hidden_layer: 定义了其内部卷积层的维度。
    • num_classes: 这是关键。新的掩码预测器需要为每一个类别都生成一个对应的掩码。因此，它的输出通道数也需要是 num_classes。模型会先判断物体是什么类别，然后选择该类别对应的通道输出作为最终的掩码。
  • 通过这个赋值操作，我们把原来为COCO数据集（91个类别）生成掩码的“旧头”，换成了为我们自己任务（num_classes个类别）生成掩码的“新头”。

第4步：返回模型

1

return model

• 功能: 返回这个经过完全定制和修改后的新模型。
• 解释: 这个返回的 model 保留了强大的预训练骨干网络，但拥有了两个为我们特定任务量身定制的“新头”（一个用于分类和边界框回归，一个用于掩码生成）。现在，你可以用这个模型在你自己的实例分割数据集上进行训练了。

总结

这个函数完美地展示了如何通过迁移学习来创建一个自定义的实例分割模型。整个过程可以概括为：

1. 加载一个通用的、预训练好的Mask R-CNN模型。
2. 更换“分类头”：使其能够识别你自定义的物体类别。
3. 更换“掩码头”：使其能够为你的每一个自定义类别生成对应的像素级掩码。

这样，你就能利用COCO数据集上学到的强大特征提取能力，高效地解决自己的实例分割问题，而无需从零开始训练整个庞大的网络。

这段代码展示了 torchvision 目标检测模型在**训练（Training）和推断（Inference）**两种模式下截然不同的行为。

我们来分两部分详细解析。

第一部分：训练模式 (For Training)

逐行解释:

1. images, targets = next(iter(data_loader)):

   • iter(data_loader) 创建一个迭代器。
   • next(...) 从迭代器中取出一项，也就是一个批次的数据。
   • images 是一个包含多张图像的张量（如果图像大小相同并通过默认collate_fn堆叠）或一个元组（如果图像大小不同，由自定义的 collate_fn 处理）。
   • targets 是一个包含了对应标注信息的元组或列表。每个元素通常是一个字典，包含 boxes, labels 等键。

2. images = list(image for image in images):

   • 关键点: torchvision 的目标检测模型在训练和推断时，都希望接收一个Python列表作为图像输入，其中列表的每个元素都是一个 [C, H, W] 形状的张量。
   • 这个操作确保了 images 是一个列表，即使 DataLoader 因为某些原因将它们打包成了单个大张量。

3. targets = [{k: v for k, v in t.items()} for t in targets]:

   • 这行代码确保 targets 是一个字典的列表。每个字典代表一张图像的真实标注（ground-truth），包含 'boxes', 'labels', 'masks' 等键值对。这通常是 DataLoader 和 collate_fn 已经处理好的标准格式，这里是做一个确认。

4. output = model(images, targets):

   • 这是最核心的区别：当模型同时接收到 images 和 targets（真实标注）时，它会自动进入训练模式。
   • 在这种模式下，模型会执行以下操作：
     1. 对输入图像进行正向传播，得到预测结果。
     2. 将预测结果与你提供的 targets (真实标注) 进行比较。
     3. 计算各种损失函数的值。
   • 因此，返回的 output 不是预测结果，而是一个包含所有损失项的字典。

5. print(output):

   • 打印出的结果会是类似这样的一个字典：

     1 2 3 4

     {'loss_classifier': tensor(0.1234, grad_fn=<...>), 'loss_box_reg': tensor(0.2345, grad_fn=<...>), 'loss_objectness': tensor(0.3456, grad_fn=<...>), 'loss_rpn_box_reg': tensor(0.4567, grad_fn=<...>)}

   • 这些损失值随后会在训练循环中用于反向传播（loss.backward()）和更新模型权重。

第二部分：推断模式 (For Inference)

逐行解释:

1. model.eval():

   • 至关重要的一步！ 这行代码将模型切换到评估模式。
   • 这会关闭 Dropout 和 BatchNorm 等在训练和推断时行为不同的层。如果不做这一步，每次预测的结果可能会因为这些层的随机性而不同，并且会使用批次数据的统计信息，导致结果不准确。

2. x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]:

   • 这里我们创建了一个虚拟的输入 x。它是一个包含两个张量的列表，模拟一个批次包含两张不同尺寸的图像。这再次强调了模型输入的格式是一个列表。

3. predictions = model(x):

   • 核心区别：当模型只接收到 images（这里是x），而没有 targets 时，它会自动进入推断模式。
   • 在这种模式下，模型只执行正向传播，并返回最终处理好的预测结果。

4. print(predictions[0]):

   • 返回的 predictions 是一个列表，其长度等于输入图像的数量（这里是2）。列表中的每个元素都是一个字典，包含了对相应输入图像的预测结果。
   • 打印 predictions[0] 会输出对第一张图像的预测，其格式会是类似这样的一个字典：

     1 2 3

     {'boxes': tensor([[x1, y1, x2, y2], ...]), 'labels': tensor([label1, ...]), 'scores': tensor([score1, ...])}

   • boxes: 检测到的边界框坐标 [N, 4]。
   • labels: 每个边界框对应的类别标签 [N]。
   • scores: 每个检测结果的置信度得分（0到1之间）[N]。通常我们会根据这个得分来过滤掉置信度低的预测结果。

总结

1
2
3

indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices[:-50])
dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])

好的，我们来详细解释这段代码。

这段代码的目的是执行机器学习中一个非常基础且重要的步骤：将一个完整的数据集随机分割成一个训练集（training set）和一个测试集（test set）。

可以把它想象成洗牌。你有一整副牌（dataset），你想留出一小部分（比如50张）来做最后的测试，剩下的牌用来练习。为了保证公平，你不能只把顶上或底下的50张拿出来，而是要先彻底洗牌，然后从洗好的牌堆里分出两部分。

让我们用一个具体的例子来解释，假设我们原始的 dataset 里有 1000 个样本（比如1000张图片和它们的标注）。

第1步：生成随机索引

1
2

# 将数据集分割成训练集和测试集
indices = torch.randperm(len(dataset)).tolist()

• len(dataset): 首先，获取数据集中样本的总数。在我们的例子中，len(dataset) 会返回 1000。
• torch.randperm(1000): 这是核心的“洗牌”操作。randperm 是 “random permutation”（随机置换）的缩写。这个函数会生成一个从 0 到 999 的整数序列，但是顺序是完全随机打乱的。
  • 例如，它可能会生成一个像 tensor([15, 980, 2, 541, ... , 33]) 这样的张量，其中 0 到 999 之间的每个数字都只出现一次。
• .tolist(): 这个方法将上面生成的 PyTorch 张量转换成一个普通的 Python 列表。
  • 所以 indices 现在是一个 Python 列表，如 [15, 980, 2, 541, ... , 33]。我们现在有了一个随机的索引顺序。

第2步：创建训练子集

1

dataset = torch.utils.data.Subset(dataset, indices[:-50])

• indices[:-50]: 这是一个 Python 的列表切片操作。它的意思是“取 indices 列表中的所有元素，除了最后50个”。
  • 在我们的例子中，这会得到一个包含 1000 - 50 = 950 个随机索引的列表。
• torch.utils.data.Subset(...): 这是一个非常有用的工具类。它接收一个原始的数据集和一个索引列表，然后创建一个“子集”。
  • 这个子集并不会在内存中复制数据，它只是一个“视图”或者说一个“包装器”，它知道只从原始 dataset 中提取 indices[:-50] 指定的那些样本。这样做非常高效。
• dataset = ...: 这里用新创建的训练子集覆盖了原来的 dataset 变量。现在，dataset 变量指向的这个对象只包含了950个用于训练的样本。

第3步：创建测试子集

1

dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])

重要提示： 这里的代码 torch.utils.data.Subset(dataset_test, ...) 很可能是一个笔误。它应该从原始的、完整的 dataset 中创建子集，而不是从一个已经存在的（可能是空的或不同的）dataset_test 中创建。

正确的、符合逻辑的代码应该是：

1
2

# 修正后的代码
dataset_test = torch.utils.data.Subset(dataset, indices[-50:])

让我们基于这个修正后的代码来解释：

• indices[-50:]: 这同样是列表切片操作，意思是“只取 indices 列表中的最后50个元素”。
• torch.utils.data.Subset(dataset, ...): 我们再次使用 Subset 类，但这次我们传入的是原始 dataset 和最后50个随机索引。
• dataset_test = ...: 这样就创建了一个名为 dataset_test 的新对象，它包含了50个专门用于测试的样本。

总结

整个过程的结果是：

1. 我们从原始的1000个样本中，随机、不重复地挑选了950个样本组成了新的 dataset（训练集）。
2. 我们用剩下的50个样本组成了 dataset_test（测试集）。

由于索引是随机打乱后分割的，这就保证了训练集和测试集是随机分配的，并且没有任何交集（即一个样本不会同时出现在训练集和测试集中）。这是进行可靠模型评估的基础，可以防止模型“泄题”（即在训练中见过了测试题），从而更公平地衡量模型的泛化能力。

代码展示了如何将一个训练好的实例分割（或目标检测）模型的推断结果（inference result）进行可视化。

整个过程可以分为三个主要步骤：

1. 获取模型预测：将一张图片喂给模型，得到预测的边界框、标签和掩码。
2. 准备可视化数据：处理原始图片和预测结果，使其符合可视化函数要求的格式。
3. 绘制并显示：使用 torchvision.utils 和 matplotlib 将边界框和掩码画在图片上并最终显示出来。

第1步：准备工作和模型推断

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

import matplotlib.pyplot as plt
from torchvision.utils import draw_bounding_boxes, draw_segmentation_masks

# 读取一张图片
image = read_image("data/PennFudanPed/PNGImages/FudanPed00046.png")
# 获取用于评估的图像变换（通常包括缩放和归一化）
eval_transform = get_transform(train=False)

# 将模型设置为评估模式
model.eval()
# 使用 no_grad() 上下文管理器，因为我们不需要计算梯度
with torch.no_grad():
    # 对原始图像应用变换
    x = eval_transform(image)
    # 将4通道的RGBA图像转换为3通道的RGB，并移动到正确的设备（CPU或GPU）
    x = x[:3, ...].to(device)
    # 将处理后的图像喂给模型（注意要包装在列表中）
    predictions = model([x, ])
    # 从预测结果列表中获取我们单个图像的预测
    pred = predictions[0]

• import ...: 导入绘图库matplotlib和torchvision提供的两个非常方便的可视化工具：draw_bounding_boxes和draw_segmentation_masks。
• read_image(...): 加载一张图片。torchvision.io.read_image 会将其读取为一个 PyTorch 张量，形状为 [Channels, Height, Width]，数值范围通常是 0-255。
• model.eval(): 非常重要。将模型切换到评估（推断）模式，这会关闭 Dropout 等层，确保每次预测结果是确定的。
• with torch.no_grad(): 告诉 PyTorch 在这个代码块里不要计算梯度。这可以显著减少内存消耗并加速计算，因为在推断时我们不需要反向传播。
• x = x[:3, ...]: 这一步是在处理图像通道。输入的PNG图片可能是4通道的RGBA（红、绿、蓝、透明度）。大多数预训练模型只接受3通道的RGB图像。这个切片操作 [:3, ...] 就是取前三个通道（R, G, B），丢弃Alpha通道。
• predictions = model([x, ]): 这是执行推断的核心。注意，x 被包装在一个列表中 [x]。这是因为torchvision的模型设计为可以接收一个批次（batch）的图像进行处理。即使只预测一张图，也要把它当作一个大小为1的批次传入。
• pred = predictions[0]: 模型返回的是一个列表，列表中每个元素对应输入批次中每张图的预测结果。因为我们只输入了一张图，所以我们取出列表的第一个元素，这就是我们需要的预测字典。pred 现在是一个包含 'boxes', 'labels', 'scores', 'masks' 等键的字典。

第2步：准备用于绘图的数据

1
2
3
4
5
6
7
8

# 将原始图像的像素值缩放到0-255范围，并转换为uint8整数类型
image = (255.0 * (image - image.min()) / (image.max() - image.min())).to(torch.uint8)
# 同样，确保用于显示的图像也是3通道的RGB
image = image[:3, ...]
# 为每个预测框创建一个带置信度得分的标签字符串
pred_labels = [f"pedestrian: {score:.3f}" for label, score in zip(pred["labels"], pred["scores"])]
# 获取预测的边界框，并转换为整数类型
pred_boxes = pred["boxes"].long()

• image = (255.0 * ...): 这是一个归一化操作，用于将image张量的像素值安全地转换到 [0, 255] 的范围内，并转换为uint8类型。这是绘图函数所期望的标准图像格式。
• pred_labels = [...]: 这是一个列表推导式，用于创建更友好的标签。它遍历预测出的每个标签和对应的置信度分数，然后格式化成一个字符串，例如 "pedestrian: 0.998"。
• pred_boxes.long(): 绘图函数要求边界框的坐标是整数像素值，所以我们用 .long() 或 .int() 将浮点数坐标转换成整数。

第3步：绘制边界框和掩码

1
2
3
4
5
6
7

# 在原始图像上绘制边界框和标签，颜色为红色
output_image = draw_bounding_boxes(image, pred_boxes, pred_labels, colors="red")

# 对模型的软掩码进行阈值处理，得到二进制掩码
masks = (pred["masks"] > 0.7).squeeze(1)
# 在已经画好边界框的图像上，再叠加上蓝色的分割掩码
output_image = draw_segmentation_masks(output_image, masks, alpha=0.5, colors="blue")

• draw_bounding_boxes(...): 第一个绘图步骤。它接收原始图像、处理好的边界框和标签，然后返回一个新的图像张量，上面已经画好了红色的边界框和文字。
• masks = (pred["masks"] > 0.7).squeeze(1): 这是一个关键的掩码预处理步骤。
  • pred["masks"]: 模型输出的通常是“软掩码”，每个像素值在0到1之间，表示该像素属于物体的概率。
  • > 0.7: 我们设置一个阈值（这里是0.7）。所有概率大于0.7的像素被认为是物体的一部分（变为True），其他的则不是（变为False）。这样我们就从软掩码得到了二进制（黑白）掩码。
  • .squeeze(1): 模型输出的掩码张量形状通常是 [N, 1, H, W]（N个物体，1个通道，高，宽）。但绘图函数期望的输入是 [N, H, W]。.squeeze(1) 会移除那个大小为1的通道维度。
• draw_segmentation_masks(...): 第二个绘图步骤。它在output_image（已经有边界框的图）的基础上，再把蓝色的掩码叠加上去。alpha=0.5 参数设置了掩码的透明度为50%，这样我们还能看到掩码下方的原始图像。

第4步：显示最终结果

1
2
3
4

# 创建一个12x12英寸的画布
plt.figure(figsize=(12, 12))
# 显示最终的图像
plt.imshow(output_image.permute(1, 2, 0))

• plt.figure(...): 设置 matplotlib 画布的大小，让最终显示的图片更大更清晰。
• plt.imshow(...): 使用 imshow 来显示图像。
• output_image.permute(1, 2, 0): 至关重要的一步。PyTorch 和 torchvision 处理图像时，默认的维度顺序是 [Channels, Height, Width] (C, H, W)。而 matplotlib 和大多数图像处理库期望的顺序是 [Height, Width, Channels] (H, W, C)。.permute(1, 2, 0) 就是在重新排列张量的维度，将顺序从 (0, 1, 2) 调整为 (1, 2, 0)，以满足 imshow 的要求。没有这一步，imshow 会报错或者显示出颜色奇怪的错误图像。