李沐简易SSD单发多框检测实现（详细注释整理）

简介

单发多框检测（SSD）
SSD (Single Shot MultiBox Detector) 是一种目标检测算法，它可以在一张图像中同时检测多个目标。
它采用单次卷积网络（single shot convolutional network）来进行检测，因此又被称为单次检测器。
SSD 通过使用预先训练的卷积神经网络来提取图像的特征，然后使用多个不同尺度的滑动窗口来检测目标。
这种方法可以同时检测不同大小的目标，并且具有较高的检测精度。、
1
2
3
4
5

13.7.1. 模型

13.7.1.1. 类别预测层

%matplotlib inline
import torch
import torchvision
from torch import nn
from torch.nn import functional as F
from d2l import torch as d2l


def cls_predictor(num_inputs, num_anchors, num_classes):
    return nn.Conv2d(num_inputs, num_anchors * (num_classes + 1),
                     kernel_size=3, padding=1)
"""
这段代码定义了一个名为cls_predictor的函数，它用于创建一个卷积层，用于对检测框进行分类。该函数接受三个参数：

num_inputs：输入通道数
num_anchors：锚框数量
num_classes：分类类别数
函数体内使用了 PyTorch 的 nn.Conv2d 模块来创建卷积层。 它的输入通道数为 num_inputs，
输出通道数为 num_anchors * (num_classes + 1)，卷积核大小为3， padding 为1。这个卷积层用于对检测框进行分类预测。

这一步是SSD中检测部分。
"""
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

13.7.1.2. 边界框预测层

def bbox_predictor(num_inputs, num_anchors):
    return nn.Conv2d(num_inputs, num_anchors * 4, kernel_size=3, padding=1)
"""
这段代码定义了一个名为bbox_predictor的函数，它用于创建一个卷积层，用于对检测框进行边界框预测。该函数接受两个参数：

num_inputs：输入通道数
num_anchors：锚框数量
函数体内使用了 PyTorch 的 nn.Conv2d 模块来创建卷积层。 它的输入通道数为 num_inputs，
输出通道数为 num_anchors * 4，卷积核大小为3， padding 为1。这个卷积层用于对检测框进行边界框预测。

这一步是SSD中边界框预测部分。
"""
1
2
3
4
5
6
7
8
9
10
11
12

13.7.1.3. 连结多尺度的预测

def forward(x, block):
    return block(x)

Y1 = forward(torch.zeros((2, 8, 20, 20)), cls_predictor(8, 5, 10))
Y2 = forward(torch.zeros((2, 16, 10, 10)), cls_predictor(16, 3, 10))
Y1.shape, Y2.shape
"""这段代码定义了一个名为 forward 的函数，它接受两个参数:

x：输入张量
block：一个由 cls_predictor 或 bbox_predictor 创建的卷积层
函数的主要作用是调用这个卷积层对输入进行卷积运算。

在这里有两个示例使用这个函数，一个使用 cls_predictor 创建卷积层，输入张量为 (2, 8, 20, 20)，
锚框数量为 5，类别数为 10，输出张量的形状为 (2, 50, 20, 20)。
另一个使用 cls_predictor 创建卷积层，输入张量为 (2, 16, 10, 10)，锚框数量为 3，类别数为 10，输出张量的形状为 (2, 30, 10, 10)。
"""
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

def flatten_pred(pred):
    return torch.flatten(pred.permute(0, 2, 3, 1), start_dim=1)

def concat_preds(preds):
    return torch.cat([flatten_pred(p) for p in preds], dim=1)
"""
这段代码定义了两个函数 flatten_pred 和 concat_preds。

flatten_pred 函数将输入预测张量进行降维处理，将预测张量的多维度拉平成一维，这样就可以方便地将多个预测张量拼接在一起。
concat_preds 函数将多个预测张量进行拼接，使用 PyTorch 的 torch.cat 函数将多个预测张量在第一维（batch size）上进行拼接。
举个例子，假设有3个预测张量，第一个预测张量的形状为 (2, 50, 20, 20)，第二个预测张量的形状为 (2, 30, 10, 10)，
第三个预测张量的形状为 (2, 40, 5, 5)。在调用 concat_preds 函数之后，得到的拼接后的张量的形状为 (2, 80, 20, 20)。

这两个函数主要是为了将多个预测张量进行拼接。
"""
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

concat_preds([Y1, Y2]).shape
"""
concat_preds 函数将多个预测张量进行拼接，在这里，我们调用了 concat_preds([Y1, Y2])，将 Y1 和 Y2 两个预测张量拼接在一起。

对于 Y1 的形状为 (2,50,20,20)，对于 Y2 的形状为 (2,30,10,10)。在拼接之后得到的拼接后的张量的形状为 (2,80,20,20),
即 Y1.shape[0] = Y2.shape[0] and Y1.shape[2:] = Y2.shape[2:].

这里第一维是batch size，第二维是80个通道，是Y1和Y2的通道数的总和，第三维和第四维是20*20，是Y1和Y2的空间维度的总和。
"""
1
2
3
4
5
6
7
8
9

13.7.1.4. 高和宽减半块

def down_sample_blk(in_channels, out_channels):
    blk = []
    for _ in range(2):
        blk.append(nn.Conv2d(in_channels, out_channels,
                             kernel_size=3, padding=1))
        blk.append(nn.BatchNorm2d(out_channels))
        blk.append(nn.ReLU())
        in_channels = out_channels
    blk.append(nn.MaxPool2d(2))
    return nn.Sequential(*blk)
"""
这段代码定义了一个名为 down_sample_blk 的函数, 这个函数用于创建一个下采样块。它接受两个参数：

in_channels：输入通道数
out_channels：输出通道数
函数体内创建了一个空的列表 blk，然后使用循环语句进行重复添加元素，其中每次添加了3个元素，
即一个卷积层，一个批量归一化层和一个ReLU层。每次循环完成后,更新in_channels的值为out_channels, 然后最后添加一个最大池化层.

这些层的组合是一个下采样块，这个下采样块可以用来缩小图像尺寸，同时增加深度。最后返回由这些层组成的序列模块。
这一步是SSD中下采样部分。
"""
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

forward(torch.zeros((2, 3, 20, 20)), down_sample_blk(3, 10)).shape
"""
这句话将输入张量 (2, 3, 20, 20) 传入名为 down_sample_blk 的函数中，并将其输出传入 forward 函数中进行处理。

down_sample_blk 函数中, in_channels = 3, out_channels = 10, 首先会在 blk 列表中添加 2 个卷积层， 2 个批量归一化层和2个ReLU层，
然后最后添加一个最大池化层。

这些层的组合是一个下采样块，这个下采样块可以用来缩小图像尺寸，同时增加深度。

输入张量的形状为 (2, 3, 20, 20)，首先通过两个卷积层和批量归一化层，每次进行卷积运算后都会增加通道数，
最后经过一个最大池化层进行下采样操作，
得到的输出张量的形状为 (2, 10, 10, 10)。

因此输出的张量的形状为 (2,10,10,10)，这个下采样块可以用来缩小图像尺寸，同时增加深度，这样在更深的层中可以更好的捕捉细节。
"""
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

13.7.1.5. 基本网络块

def base_net():
    blk = []
    num_filters = [3, 16, 32, 64]
    for i in range(len(num_filters) - 1):
        blk.append(down_sample_blk(num_filters[i], num_filters[i+1]))
    return nn.Sequential(*blk)

forward(torch.zeros((2, 3, 256, 256)), base_net()).shape
"""
这段代码定义了一个名为 base_net 的函数，该函数用于创建一个基础网络。在这里，函数内部创建了一个空列表 blk。

num_filters = [3, 16, 32, 64], 用于存储输入和输出通道数,然后使用循环语句进行重复添加元素。其中每次添加一个下采样块，
并使用num_filters[i]作为输入通道数，num_filters[i+1]作为输出通道数。

输入张量的形状为 (2, 3, 256, 256)， 通过循环添加的下采样块，在经过多次下采样操作后得到的输出张量的形状为 (2, 64, 32, 32)。

所以 base_net 函数返回的是这些下采样块组成的序列模块, 输入一个大尺寸图像，经过多次下采样操作后尺寸变小，通道数变多，
这样在更深的层中可以更好的捕捉细节。
"""
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

13.7.1.6. 完整的模型

def get_blk(i):
    if i == 0:
        blk = base_net()
    elif i == 1:
        blk = down_sample_blk(64, 128)
    elif i == 4:
        blk = nn.AdaptiveMaxPool2d((1,1))
    else:
        blk = down_sample_blk(128, 128)
    return blk
"""
这段代码定义了一个名为 get_blk 的函数，该函数用于根据输入的索引 i 返回不同的模块。

函数中使用了 if-elif 语句来控制返回不同的模块：

当 i 为 0 时，返回基础网络 base_net()。
当 i 为 1 时，返回下采样块 down_sample_blk(64, 128)。
当 i 为 4 时，返回自适应最大池化层 nn.AdaptiveMaxPool2d((1,1))。
其余情况下，返回下采样块 down_sample_
"""
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

def blk_forward(X, blk, size, ratio, cls_predictor, bbox_predictor):
    Y = blk(X)
    anchors = d2l.multibox_prior(Y, sizes=size, ratios=ratio)
    cls_preds = cls_predictor(Y)
    bbox_preds = bbox_predictor(Y)
    return (Y, anchors, cls_preds, bbox_preds)
"""
这段代码定义了一个名为blk_forward的函数,这个函数用于在给定模块和参数的情况下进行前向计算。

函数接受五个参数：

X ：输入数据
blk ：需要使用的模块
size ：锚框大小
ratio ：锚框宽高比
cls_predictor ：分类预测器
bbox_predictor ：边界框预测器
首先，使用传入的模块对输入数据进行前向计算，得到输出 Y。然后，使用 d2l.multibox_prior 函数计算锚框，传入输出 Y 、
锚框大小和锚框宽高比来生成锚框。之后,使用 cls_predictor 和 bbox_predictor 对 Y 进行分类预测和边界框预测,
得到 cls_preds 和 bbox_preds。

最后,函数将得到的 Y, anchors, cls_preds, bbox_preds 作为元组返回.这个函数可以用来处理不同层的预测结果，并得到最终的预测结果。
"""
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

sizes = [[0.2, 0.272], [0.37, 0.447], [0.54, 0.619], [0.71, 0.79],
         [0.88, 0.961]]
ratios = [[1, 2, 0.5]] * 5
num_anchors = len(sizes[0]) + len(ratios[0]) - 1
"""
这段代码中定义了两个列表 sizes 和 ratios。

sizes 列表中存储了不同层中锚框的大小，每个子列表中存储了该层中锚框的高和宽。
ratios 列表中存储了不同层中锚框的宽高比，每个子列表中存储了该层中锚框的宽高比。
然后, num_anchors 被赋值为总锚框数，其等于每层的锚框大小数量与宽高比数量之和，其中每层减1.
这个sizes和ratios是用来生成不同层的锚框的。
"""
1
2
3
4
5
6
7
8
9
10
11
12

class TinySSD(nn.Module):
    def __init__(self, num_classes, **kwargs):
        super(TinySSD, self).__init__(**kwargs)
        self.num_classes = num_classes
        idx_to_in_channels = [64, 128, 128, 128, 128]
        for i in range(5):
            # 即赋值语句self.blk_i=get_blk(i)
            setattr(self, f'blk_{i}', get_blk(i))
            setattr(self, f'cls_{i}', cls_predictor(idx_to_in_channels[i],
                                                    num_anchors, num_classes))
            setattr(self, f'bbox_{i}', bbox_predictor(idx_to_in_channels[i],
                                                      num_anchors))

    def forward(self, X):
        anchors, cls_preds, bbox_preds = [None] * 5, [None] * 5, [None] * 5
        for i in range(5):
            # getattr(self,'blk_%d'%i)即访问self.blk_i
            X, anchors[i], cls_preds[i], bbox_preds[i] = blk_forward(
                X, getattr(self, f'blk_{i}'), sizes[i], ratios[i],
                getattr(self, f'cls_{i}'), getattr(self, f'bbox_{i}'))
        anchors = torch.cat(anchors, dim=1)
        cls_preds = concat_preds(cls_preds)
        cls_preds = cls_preds.reshape(
            cls_preds.shape[0], -1, self.num_classes + 1)
        bbox_preds = concat_preds(bbox_preds)
        return anchors, cls_preds, bbox_preds
"""
    这段代码定义了一个名为 TinySSD 的类，继承了 PyTorch 的 nn.Module 类。

在 TinySSD 类的构造函数中，首先将输入的类别数量赋值给 self.num_classes。然后，定义了一个名为 idx_to_in_channels 的列表，
用于存储不同层的输入通道数。

然后,使用循环语句对于每一层进行如下操作:

通过 get_blk(i) 函数获取第 i 个模块并将其赋值给 self.blk_i 。
通过 cls_predictor函数获取第 i 个分类预测器并将其赋值给 self.cls_i 。

通过 bbox_predictor 函数获取第 i 个边界框预测器并将其赋值给 self.bbox_i 。
在 forward 函数中，对于每一层使用 blk_forward 函数进行前向计算，得到锚框、分类预测和边界框预测。
最后，将所有层的锚框、分类预测和边界框预测拼接在一起并返回。

TinySSD 类包含了一个完整的单发多框检测网络，可以对图像进行预测，并得到锚框、分类预测和边界框预测。
这些预测结果可以用来计算置信度和坐标，以确定最终的目标检测结果。

在使用TinySSD类进行预测之前,需要对该类进行训练。训练时需要准备大量标记数据,来训练分类预测和边界框预测.
"""
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46

net = TinySSD(num_classes=1)
X = torch.zeros((32, 3, 256, 256))
anchors, cls_preds, bbox_preds = net(X)

print('output anchors:', anchors.shape)
print('output class preds:', cls_preds.shape)
print('output bbox preds:', bbox_preds.shape)
"""
这段代码创建了一个 TinySSD 网络的实例，并将其赋值给变量 net。然后创建一个大小为 (32, 3, 256, 256) 的输入张量 X，
并使用 net 进行前向计算。

前向计算返回了三个张量：锚框、分类预测和边界框预测。最后，打印了这三个张量的形状，
分别是：(32, 5444, 4)、(32, 21776, 2) 和 (32, 21776, 4)。这些张量的大小取决于输入数据的大小和模型的设计.
"""
1
2
3
4
5
6
7
8
9
10
11
12
13
14

output anchors: torch.Size([1, 5444, 4])
output class preds: torch.Size([32, 5444, 2])
output bbox preds: torch.Size([32, 21776])
1
2
3

13.7.2. 训练模型

13.7.2.1. 读取数据集和初始化

batch_size = 32
train_iter, _ = d2l.load_data_bananas(batch_size)
"""
这段代码使用d2l库中的函数加载了一个名为 bananas 的数据集。其中,train_iter是一个迭代器，每次迭代返回一个batch大小的数据。
这些数据用于训练TinySSD网络。
"""
1
2
3
4
5
6

read 1000 training examples
read 100 validation examples
1
2

device, net = d2l.try_gpu(), TinySSD(num_classes=1)
trainer = torch.optim.SGD(net.parameters(), lr=0.2, weight_decay=5e-4)
"""这段代码将TinySSD网络的实例转移到 GPU 上（如果可用的话），并创建了一个新的优化器。其中，

device 是一个 PyTorch tensor 类型的变量，表示当前使用的设备（CPU 或 GPU）。
net 是 TinySSD 网络的实例。
trainer 是一个 PyTorch 优化器，使用随机梯度下降 (SGD) 算法并对 net 的参数进行优化，其学习率为 0.2，权值衰减系数为 5e-4.
此后,可以使用训练数据和优化器训练网络,并使用验证数据进行评估,以提高模型性能.
"""
1
2
3
4
5
6
7
8
9

13.7.2.2. 定义损失函数和评价函数

cls_loss = nn.CrossEntropyLoss(reduction='none')
bbox_loss = nn.L1Loss(reduction='none')

def calc_loss(cls_preds, cls_labels, bbox_preds, bbox_labels, bbox_masks):
    batch_size, num_classes = cls_preds.shape[0], cls_preds.shape[2]
    cls = cls_loss(cls_preds.reshape(-1, num_classes),
                   cls_labels.reshape(-1)).reshape(batch_size, -1).mean(dim=1)
    bbox = bbox_loss(bbox_preds * bbox_masks,
                     bbox_labels * bbox_masks).mean(dim=1)
    return cls + bbox
"""
这段代码定义了两个损失函数：

cls_loss是交叉熵损失函数，用于计算类别预测和标签之间的差距。
bbox_loss是 L1 损失函数，用于计算边界框预测和标签之间的差距。
然后定义了 calc_loss 函数,用于计算类别预测和边界框预测之间的差距，并返回总损失。

其中，cls_preds 是类别预测，cls_labels 是类别标签，bbox_preds 是边界框预测，
bbox_labels 是边界框标签，bbox_masks 是边界框掩码。
其中 cls_loss(cls_preds.reshape(-1, num_classes),cls_labels.reshape(-1)) 计算了每个样本的交叉熵损失。
reshape(batch_size, -1) 得到了每个样本的所有锚框的总损失和,mean(dim=1)得到了每个样本的平均损失。
bbox_loss(bbox_preds * bbox_masks,bbox_labels * bbox_masks) 计算了每个样本的边界框的 L1 损失，
其中 bbox_masks 用于掩码掉无效的锚框。

mean(dim=1)得到了每个样本的平均边界框损失。
最终，cls 和 bbox 分别是每个样本的类别损失和边界框损失，它们相加得到了总损失。
"""
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

def cls_eval(cls_preds, cls_labels):
    # 由于类别预测结果放在最后一维，argmax需要指定最后一维。
    return float((cls_preds.argmax(dim=-1).type(
        cls_labels.dtype) == cls_labels).sum())

def bbox_eval(bbox_preds, bbox_labels, bbox_masks):
    return float((torch.abs((bbox_labels - bbox_preds) * bbox_masks)).sum())
"""
这段代码定义了两个评估函数，用于在训练期间评估类别预测和边界框预测的准确性：

cls_eval 函数使用 PyTorch 的 argmax 函数计算每个预测的类别与标签之间的匹配数。它返回一个浮点数，表示正确预测的数量。
bbox_eval 函数使用 PyTorch 的 abs 函数计算预测边界框和标签之间的绝对差值。它返回一个浮点数，表示所有边界框的绝对误差之和。
它们将在训练时使用,来评估类别预测和边界框预测的准确性,并通过这两个函数来监控训练过程的进度。
"""
1
2
3
4
5
6
7
8
9
10
11
12
13
14

13.7.2.3. 训练模型

num_epochs, timer = 20, d2l.Timer()
animator = d2l.Animator(xlabel='epoch', xlim=[1, num_epochs],
                        legend=['class error', 'bbox mae'])
net = net.to(device)
for epoch in range(num_epochs):
    # 训练精确度的和，训练精确度的和中的示例数
    # 绝对误差的和，绝对误差的和中的示例数
    metric = d2l.Accumulator(4)
    net.train()
    for features, target in train_iter:
        timer.start()
        trainer.zero_grad()
        X, Y = features.to(device), target.to(device)
        # 生成多尺度的锚框，为每个锚框预测类别和偏移量
        anchors, cls_preds, bbox_preds = net(X)
        # 为每个锚框标注类别和偏移量
        bbox_labels, bbox_masks, cls_labels = d2l.multibox_target(anchors, Y)
        # 根据类别和偏移量的预测和标注值计算损失函数
        l = calc_loss(cls_preds, cls_labels, bbox_preds, bbox_labels,
                      bbox_masks)
        l.mean().backward()
        trainer.step()
        metric.add(cls_eval(cls_preds, cls_labels), cls_labels.numel(),
                   bbox_eval(bbox_preds, bbox_labels, bbox_masks),
                   bbox_labels.numel())
    cls_err, bbox_mae = 1 - metric[0] / metric[1], metric[2] / metric[3]
    animator.add(epoch + 1, (cls_err, bbox_mae))
print(f'class err {cls_err:.2e}, bbox mae {bbox_mae:.2e}')
print(f'{len(train_iter.dataset) / timer.stop():.1f} examples/sec on '
      f'{str(device)}')
"""
上述代码是训练TinySSD模型的代码，它包括以下几个部分：

循环每个epoch:
定义一个Accumulator来累加训练精度的和，训练精度的和中的示例数，绝对误差的和，绝对误差的和中的示例数
将网络设置为训练模式
循环每个batch:
将梯度清零
将输入数据和标签移动到GPU
前向传播，生成多尺度的锚框，为每个锚框预测类别和偏移量
为每个锚框标注类别和偏移量
根据类别和偏移量的预测和标注值计算损失函数
反向传播，计算梯度并更新网络参数
累加训练精度的和，训练精度的和中的示例数，绝对误差的和，绝对误差的和中的示例数
计算类别错误率和偏移量绝对平均误差
将这些指标添加到动画器中
输出最终的类别错误率和偏移量绝对平均误差
输出每秒处理的样本数

首先使用训练数据进行训练，每个epoch都需要对所有训练数据进行前向传播和反向传播。
对于每个样本，都需要计算预测的类别和偏移量，并与标签进行比较以计算损失函数。在每个epoch结束后，
会计算并记录分类误差和坐标误差，并将结果加入动画中。最后，使用计时器计算每秒处理样本数。
"""
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53

class err 3.33e-03, bbox mae 3.34e-03
7872.2 examples/sec on cuda:0
1
2

13.7.3. 预测目标

X = torchvision.io.read_image(r'D:\pytorch_Pycharm\imgs\banana2.png').unsqueeze(0).float()
img = X.squeeze(0).permute(1, 2, 0).long()
"""
这段代码读取了一张图片“banana2.jpg”，然后将其转化为一个4维的张量（batch_size, channel, height, width)。
将第一维的batch size设置为1，然后将其转换为3维张量(channel, height, width)，
最后使用permute函数将其转换为(height, width, channel)并转换为long类型。
"""
1
2
3
4
5
6
7

def predict(X):
    net.eval()
    anchors, cls_preds, bbox_preds = net(X.to(device))
    cls_probs = F.softmax(cls_preds, dim=2).permute(0, 2, 1)
    output = d2l.multibox_detection(cls_probs, bbox_preds, anchors)
    idx = [i for i, row in enumerate(output[0]) if row[0] != -1]
    return output[0, idx]

output = predict(X)
"""
这段代码中，首先设置网络为评估模式，然后将输入图片X传入网络中并得到输出anchors, cls_preds, bbox_preds。

接着使用softmax函数对cls_preds进行归一化，使其成为概率值。由于softmax函数默认在最后一维进行softmax，
而在这里是在第二维，所以使用permute函数将其转换为(batch_size, num_anchors, num_classes)

最后，使用multibox_detection函数将预测的类别概率和偏移量转换为边界框。由于预测的类别可能为背景，
所以还需要将其过滤掉，即选出输出中类别不为背景的边界框。

我们使用列表推导式从输出结果中提取出所有预测类别不为-1的行的索引。
对于每一行的输出，第一个元素是预测的类别，如果预测类别为-1，表示当前行不是有效预测结果。
因此，我们只保留预测类别不为-1的行，这些行才是有效预测结果。
"""
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

def display(img, output, threshold):
    d2l.set_figsize((5, 5))
    fig = d2l.plt.imshow(img)
    for row in output:
        score = float(row[1])
        if score < threshold:
            continue
        h, w = img.shape[0:2]
        bbox = [row[2:6] * torch.tensor((w, h, w, h), device=row.device)]
        d2l.show_bboxes(fig.axes, bbox, '%.2f' % score, 'w')

display(img, output.cpu(), threshold=0.4)
"""
这段代码是将图片和预测结果显示在一起。其中，img是原始图片，output是模型预测出的结果，
threshold是过滤掉置信度过低的预测结果的阈值。在显示时，会循环遍历output中的每一行，对于每一行的置信度都要比较是否大于阈值。
如果置信度大于阈值，则显示该预测结果的边界框。
"""
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17