12｜模型预测：让模型告诉我们组件的信息

你好，我是柳博文，欢迎和我一起学习前端工程师的AI实战课。

上一节课我们完成了模型的训练，并通过分析训练结果可以知道，模型达到了不错的收敛效果。那么，这节课我们就来一起看看如何使用训练好的模型来识别和定位。

在这个过程中，我们仍然需要使用到上节课中讲述的YOLOv5的源码，源码中实现了用于模型预测的detect.py文件，所以我们像调用tarin.py文件一样调用它即可。

在进行模型检测调用前，有一些重要文件需要提前了解，那就是模型的权重文件。

模型权重文件

那么，什么是模型权重文件呢？

模型权重文件是存储神经网络模型中所有可训练参数的文件。这些参数包括神经网络中的权重和偏置，决定了模型的行为和性能。权重文件通常是在模型经过训练后保存的，它们包含了模型在训练过程中学到的知识。

我们以函数为例说明一下，一个一元一次函数 f(x) = ax+b，对于这个一元一次函数来说，如果具有一个固定的a、b值，那么给定一个输入值x，我们就会得到这个输入值x对应的固定的输出值。

那么，模型也一样，权重文件中保留了这个模型完成训练后的所有参数，对于一个视觉模型，这个输入值x就是一张图像，通过结合权重文件中的保留参数等信息进行计算，输出值y就是这张图像中对应对象的预测结果了。

所以，通常我们会说训练模型也就是在拟合一个函数，这个函数包含的参数非常多。当函数拟合到我们想要的状态时，保留下来的参数组合也就是这个权重文件。

在widgets文件夹下的权重文件是以pt为后缀名保存的，这是因为我们使用的是PyTorch来完成这个模型的训练，在PyTorch中，.pt文件是用来存储模型的权重、架构或其他信息的文件格式，它的全称是“PyTorch”。

上节课里我们训练模型后，分析了runs/train/exp/文件夹下的训练结果图。在这个文件夹下面还有一个widgets文件夹，其中存放了两个后缀名为pt的文件，分别是 last.pt 和 best.pt。这其实就是训练完成后的模型权重文件，用于模型预测的参数组合。

通过命名其实可知，一个是模型完成所有训练批次后的权重文件，另一个则是模型在训练过程中达到最优训练结果的权重文件。

完成所有训练批次后得到的权重文件比较好理解。最优训练结果的权重文件则是在训练过程中，模型根据某个指标（如mAP，mean Average Precision）在验证集上判断模型的表现所得。best.pt文件会在模型在验证集上达到最佳表现时保存下来，因此这个文件包含了最优的模型参数，通常可以用于最终的模型部署和推理。

那么，有了最优的模型权重文件，我们就可以用来进行模型预测了。

模型预测

在完成 YOLOv5 模型的训练后，我们接下来要进行模型的预测操作。这一步骤需要我们对 YOLOv5 的 detect.py 文件稍加修改，以确保预测过程的顺利进行。与修改 train.py 文件时类似，我们需要调整一些超参数以适应当前的需求。

在 YOLOv5 的源码中，detect.py 文件用于执行检测任务。首先，我们需要关注第 238 行和 239 行的代码。

    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'best.pt', help='model path or triton URL')
    parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob/screen/0(webcam)')

首先来看第一行代码，作用是设置权重文件。–weights 参数指定了要使用的模型权重文件，这里我们选择了训练过程中保存的最优模型 best.pt。这个文件保存了模型在验证集上表现最佳的权重参数，用于推理时能得到较好的预测结果。

其次是设置预测图像来源，–source 参数指定了输入数据的路径。在这个示例中，我们将需要检测的图像存放在 data/images 文件夹下。YOLOv5 支持对多个图像文件同时进行检测，只需将这些文件放在指定文件夹中即可。

接下来，我们准备两张用于检测测试的图像。图像的来源可以是数据集中用于测试的图像，也可以是与训练目标一致的新的图像。我在课程的源代码里也提供了测试图像，你可以通过这里获取。

我们把测试图像放入 data/images目录下，执行detect.py。

python detect.py

如果你的系统没有 GPU，或你希望使用 CPU 进行检测，可以使用以下命令指定设备为 CPU。

python detect.py --device cpu

YOLOv5 在进行推理时默认会使用 GPU，这是因为 GPU 在处理并行计算时具有显著的速度优势，特别是对于大规模数据和复杂场景。相比之下，CPU 虽然也能完成检测任务，但在处理速度上明显逊色于 GPU，因此检测时间会更长。

检测任务执行完成后，YOLOv5 会在 runs/detect 文件夹下生成一个新的子文件夹（例如 exp），检测结果将会放在这个文件夹中。

对于可视化图像结果，你可以直接查看这些图像，其中每个被检测到的组件都会被框选出来，并标注了类别和置信度。这使得我们可以直观地看到模型在图像中识别出的对象及其位置信息。

除了可视化图像，YOLOv5 还会生成对应的标签文件，其中记录了每个检测到的目标的类别、坐标位置和大小。通过这些数据，我们可以进一步处理，比如自动生成代码来实现界面布局。

我们一起看看后面的例子，这样更直观。首先是可视化的一个检测结果，其中左侧为原图，右侧为检测结果图。可以看到 ，我们的模型/AI组件识别助手的识别效果还不错，两张图像里的组件类型和定位都识别得比较准确，也打上了对应的矩形框和类型分类。

同时，在这个可视化检测结果下生成了一个labels文件夹。这个文件夹下以文本形式存放着检测结果，这就可以用于进行工程链路上的实现了。我们来看一下其中一个文本文件的内容构成。

0 636.5 259 1273 158
1 642.5 684 1273 544
2 638.5 1518 1277 456
2 635.5 1982.5 1269 463
2 643.5 2445 1255 452

在这个label文本文件中一共有五列，从左到右分别对应了一个属性值，class_id、x_center、y_center、width、height。

• <class_id>对应着目标的类别ID。例如，0表示第一个目标的类别，1表示第二个目标的类别，2表示接下来的目标类别。

• <x_center>则是表示目标边界框的中心点的横坐标。例如，636.5表示第一个目标边界框中心点的横坐标。

• <y_center>则是目标边界框的中心点的纵坐标。例如，259表示第一个目标边界框中心点的纵坐标。

• 代表目标边界框的宽度。例如，1273表示第一个目标的边界框的宽度。

• 代表目标边界框的高度。例如，158表示第一个目标的边界框的高度。

那么，通过这五个关键信息，我们就能够知道位于页面的什么位置，是一个什么样的组件了。再结合组件库，就能够将页面的布局进行一个简单的自动出码还原。

总结

今天的课程告一段落，我们来对这节课做一个总结吧。

在这节课中，我们使用训练好的模型进行了图像检测。首先，我们理解了模型训练完成后的权重文件。我们可以将模型看成一个包含大量参数的函数，而权重文件则是存储这些参数的地方。后缀为 .pt 的权重文件则是 PyTorch 格式的权重文件。

接着，我们使用最优权重文件进行了预测。结果显示，模型能够准确识别测试图像中的组件类型和位置，并生成了包含组件类型、位置和宽高信息的标签文件。这些关键信息可以用于实现自动化代码生成。建议你课后跟着这节课的思路动手练习一下，这样学习效果会更好。

思考题

在真实业务场景里，页面往往不止4个类别这么简单，层级结构更是复杂，那么要如何使用AI模型来完成页面布局呢？

欢迎你在留言区和我交流互动，如果这节课对你有启发，也推荐分享给身边更多朋友。