用 YOLO 识别桌球位置

发表于 2024-01-16 更新于 2024-04-26 分类于默认本文字数： 3.8k 阅读时长 ≈ 3 分钟

# 准备桌球图片 100 张

因为是固定机位，所以可以认为视角是相同的
从视频里拆出图片，最好是每隔 10-20 秒取一张图片，这样的间隔时间是因为低于这个时间图片可能还没产生变化

# 利用 LabelImg 工具，对图片进行标记，并导出为 yolo 格式

安装 LabelImg：首先，从 GitHub 仓库克隆 LabelImg 并安装其依赖项。在命令行中运行以下命令：

1
2
3

git clone https://github.com/HumanSignal/labelImg.git
cd labelImg
pip install PyQt5 lxml

运行 LabelImg：在 labelImg 目录中，运行以下命令启动 LabelImg：

1	python labelImg.py

这将打开 LabelImg 的图形界面。

打开图片：点击 “Open Dir” 按钮，然后选择包含桌球图片的文件夹。图片将显示在 LabelImg 的主窗口中。
设置标注格式：点击 “PascalVOC” 按钮，将其切换为 “YOLO”。这将使 LabelImg 以 YOLO 格式保存标注信息。
创建类别列表：点击 “Change Save Dir” 按钮，选择一个文件夹以保存标注文件。然后，创建一个名为 “classes.txt” 的文件，其中包含你要标注的所有类别名称，每个类别占一行。在本例中，你只需要一个类别，即 “table_tennis_ball”。将此文件保存在与标注文件相同的文件夹中。
标记桌球：点击 “Create RectBox” 按钮，然后在图片中的桌球周围绘制一个矩形边界框。在弹出的对话框中，输入类别名称 “table_tennis_ball”，然后按 Enter 键。如果图片中有多个桌球，请为每个桌球重复此步骤。
保存标注信息：点击 “Save” 按钮，将标注信息保存为 YOLO 格式的文本文件。文件名将与对应的图片文件名相同，但扩展名为 “.txt”。
标注更多图片：使用左右箭头按钮在图片之间导航，并重复步骤 6 和 7，为每张图片标注桌球。

完成标注后，你将得到一个包含 YOLO 格式标注信息的文本文件，可以用于训练 YOLOv5 模型。确保训练数据集中包含足够多样化的样本，以便模型能够学会准确检测桌球。

# 使用自定义数据集训练 YOLOV5 模型

# 克隆YOLOv5仓库
git clone https://github.com/ultralytics/yolov5.git

# 进入仓库目录
cd yolov5

# 安装依赖项
pip install -r requirements.txt

# 训练模型
python train.py --img 640 --batch 16 --epochs 100 --data custom_dataset.yaml --weights yolov5s.pt --cache

在这个示例中，我们使用 YOLOv5s 模型作为预训练权重， custom_dataset.yaml 文件包含了自定义数据集的相关信息。需要根据实际情况调整训练参数。

在 YOLOv5 中，训练命令的参数有多个，以下是一些常用参数的介绍：

–img：输入图片的大小。这个参数将图片调整为正方形，例如 640x640 像素。较大的图片尺寸可能会提高检测精度，但会增加计算成本和训练时间。

–batch：每个批次的图片数量。较大的批次可以更好地利用 GPU 资源，但可能会导致内存不足。

–epochs：训练周期。表示模型在整个数据集上训练的次数。较多的周期可能会提高模型性能，但会增加训练时间。

–data：数据集配置文件。这个文件包含了训练和验证数据集的路径、类别名称等信息。你需要根据自己的数据集创建一个 YAML 文件。

–weights：预训练权重文件。这个参数用于加载预训练的权重，以便进行迁移学习。你可以使用 YOLOv5 提供的预训练权重（例如 yolov5s.pt），也可以使用自己训练过的权重。

–cache：缓存数据集图片。这个参数可以加速训练，因为它会将数据集图片加载到内存中。但是，这会消耗更多的内存。

–name：实验名称。这个参数可以为你的训练实验指定一个名称，以便更好地组织和查找训练结果。

–device：训练设备。这个参数可以指定训练使用的设备，例如 "0" 表示使用第一个 GPU，"0,1" 表示使用前两个 GPU，"cpu" 表示使用 CPU。

–workers：数据加载器线程数。这个参数可以指定用于加载数据的线程数。较多的线程可以加速数据加载，但可能会占用更多的系统资源。

这些参数仅是训练命令中的一部分。可以在 YOLOv5 的官方文档中查看更多参数和详细信息。根据实际需求调整参数以优化模型性能和训练速度。

训练模型所用到的图片应该存放在一个目录中，该目录通常包含两个子目录：一个用于训练数据，另一个用于验证数据。例如，你可以按照以下结构组织你的图片数据：

dataset/
├── train/
│ ├── image1.jpg
│ ├── image2.jpg
│ └── …
└── valid/
├── image101.jpg
├── image102.jpg
└── …
此外，你还需要为每张图片提供标注信息。YOLOv5 使用一种特定的文本格式来存储标注信息。对于每张图片，你需要创建一个同名的文本文件，其中包含对象的类别和边界框坐标。例如，对于 image1.jpg，你需要创建一个 image1.txt 文件，其中包含如下信息：

0 0.5 0.5 0.25 0.25
这表示图片中有一个对象，类别为 0（例如桌球），边界框的中心点坐标为（0.5, 0.5），宽度和高度分别为 0.25。坐标和尺寸都是相对于图片宽度和高度的比例值。

一旦你准备好了图片和标注数据，你需要创建一个 YAML 配置文件，其中包含训练和验证数据集的路径、类别信息等。例如，你可以创建一个名为 custom_dataset.yaml 的文件，内容如下：

train: dataset/train/
val: dataset/valid/
nc: 1
names: [‘table_tennis_ball’]
这个配置文件指定了训练和验证数据集的路径，以及类别数量（nc）和类别名称（names）。

标注信息文件应与对应的图像文件放在同一个目录中。在我们之前提到的数据集结构示例中，标注文件应放在 train 和 valid 子目录中。每个图像文件都需要一个对应的标注文件。例如，如果你有一个名为 image1.jpg 的图像文件，你需要在同一个目录中创建一个名为 image1.txt 的标注文件。

以下是一个包含标注文件的数据集结构示例：

dataset/
├── train/
│ ├── image1.jpg
│ ├── image1.txt
│ ├── image2.jpg
│ ├── image2.txt
│ └── …
└── valid/
├── image101.jpg
├── image101.txt
├── image102.jpg
├── image102.txt
└── …
在 YOLOv5 中，标注文件使用一种特定的文本格式。每行代表一个对象，包含类别 ID、边界框中心点的 x 和 y 坐标，以及边界框的宽度和高度。所有坐标和尺寸都是相对于图像宽度和高度的比例值。例如，一个标注文件的内容可能如下：

0 0.5 0.5 0.25 0.25
这表示图像中有一个对象，类别 ID 为 0（例如桌球），边界框的中心点坐标为（0.5, 0.5），宽度和高度分别为 0.25。

准备好图像和标注文件后，你需要创建一个 YAML 配置文件，指定训练和验证数据集的路径、类别数量和类别名称。在训练命令中，使用–data 参数指定 YAML 配置文件的路径。YOLOv5 会自动读取图像和对应的标注文件，进行模型训练。

如果一张图片中包含多个桌球，你可以在标注文件中为每个桌球添加一行。每行表示一个对象，包含类别 ID、边界框中心点的 x 和 y 坐标，以及边界框的宽度和高度。所有坐标和尺寸都是相对于图像宽度和高度的比例值。

例如，如果一张图片中有两个桌球，标注文件的内容可能如下：

0 0.3 0.3 0.2 0.2
0 0.7 0.7 0.2 0.2
第一行表示第一个桌球，类别 ID 为 0，边界框的中心点坐标为（0.3, 0.3），宽度和高度分别为 0.2。第二行表示第二个桌球，类别 ID 为 0，边界框的中心点坐标为（0.7, 0.7），宽度和高度分别为 0.2。

在创建了包含多个桌球的标注文件后，YOLOv5 会自动处理这些标注信息并在训练过程中识别多个对象。确保你的数据集中包含足够的多对象样本，以便模型能够学会准确检测多个桌球。

# 使用训练好的模型进行对象检测：

import torch
from PIL import Image
from yolov5 import YOLOv5

# 加载训练好的YOLOv5模型
model = YOLOv5("runs/train/exp/weights/best.pt", device="cuda" if torch.cuda.is_available() else "cpu")

# 读取图片
image = Image.open("table_tennis_image.jpg")

# 进行对象检测
results = model.predict(image)

# 显示检测结果
results.show()