自定义训练的 YOLOv8 模型进行邮票整理

还记得人们过去常常寄信和明信片的日子吗？明信片上通常描绘了你所访问的城市或国家的一些风景如画的景色，但对许多人来说，邮票才是主要的吸引力。它们被视为微型艺术品，被认为是非常值得收藏的。

由于每张明信片、信封、包裹和邮包都需要邮票，许多人最终拥有了大量的邮票收藏。这变成了一个计算机视觉项目：我没有一张一张地拍摄每张邮票的照片，而是拍摄了邮票相册的每一页，并使用目标检测来找到并存储图片中所有邮票的图像。

目标检测

由于我之前没有尝试过目标检测（只做过图像分类），我花了一些时间寻找完成这项任务的最佳方法。我在Roboflow上发现了一个类似的项目[1]，并能够在线测试一些我自己的图像，使用已经在一个邮票数据集上训练过的YOLOv8模型。然而，我无法让模型在我自己的机器上本地运行，所以我决定使用Ultralytics的YOLOv8模型训练自己的模型。下图展示了我自定义训练的YOLOv8目标检测模型如何处理一张包含57张意大利邮票的图片。

自定义训练的 YOLOv8 模型进行邮票整理

使用自定义训练的YOLOv8模型进行邮票检测

每张检测到的邮票也自动保存为单独的裁剪图像文件，其中一些如下所示。

自定义训练的 YOLOv8 模型进行邮票整理

自动裁剪并保存的单个邮票图像

在自定义数据集上训练YOLO

加载一个YOLO模型（使用COCO数据集的预训练权重）并在你自己的机器上使用自定义数据集进行训练并不困难。一旦创建了虚拟环境并安装了所需的Python库（例如pytorch、ultralytics），实际的训练和推理只需要几行代码。

困难的部分可能是获取足够大的标注数据集。在我的情况下，我很幸运地在网上找到了一个已经标注并已转换为YOLOv8格式的邮票数据集[1]。如果你必须制作自己的数据集来让YOLO检测不属于其开箱即用的80个类别之一的东西，Roboflow和CVAT提供了注释工具，你可以使用它们来完成这项工作。

一旦数据集准备好并保存在本地文件夹中，可以使用config.yaml文件提供模型训练、验证和测试所需的类标签和路径。我的config.yaml文件如下所示：

复制

path: /home/username/venv_folder/venv_name/yolov8-project/ # absolute path to dataset
test: test/images # relative path to test images
train: train/images # relative path to training images
val: val/images # relative path to validation images

# classes
names:
  0: postage-stamp

使用自定义数据集训练YOLOv8所需的Python代码如下所示。建议使用GPU，并在训练前检查其可用性。

复制

import torch, ultralytics


# Check library versions
print("PyTorch version:", torch.__version__)
print("Ultralytics version:", ultralytics.__version__)


# Check if GPU is available
if torch.cuda.is_available():
    print(f"GPU is available: {torch.cuda.get_device_name(0)}")
else:
    print("GPU is not available")

自定义训练的 YOLOv8 模型进行邮票整理

在训练之前，加载带有预训练权重的YOLOv8模型。模型名称末尾的字母（yolov8n.pt）可用于选择模型大小。字母n表示最小且最快的模型，具有最少的可训练参数，但如果需要更高的精度且速度不是关键，你可以选择另一个模型（n,s,m,l,x）[2]。

复制

from ultralytics import YOLO


# Load model with pretrained weights (recommended)
model = YOLO("yolov8n.pt") 


# Optionally load a model without pretrained weights
# model = YOLO("yolov8n.yaml")


# You can also try yolo11n
# model = YOLO("yolo11n.pt")


# Train model
model.train(data="config.yaml", epochs=100, patience=10)

训练完成后，会生成一个runs/detect/train/文件夹，其中包含许多有用的信息，如训练/验证指标和混淆矩阵，因此请务必查看它们以更好地了解模型性能。

自定义训练的 YOLOv8 模型进行邮票整理

YOLOv8训练指标

要使用自定义训练的模型从新图像中检测对象，请从runs/detect/train/weights文件夹加载它。你可以选择best.pt或last.pt，前者是得分最高的模型，后者是最后一个epoch训练的模型。

复制

# Load the trained YOLO model
model = YOLO("/home/username/venv_folder/venv_name/runs/detect/train/weights/best.pt")


# Specify folder containing images for object detection
image_folder = "/home/username/venv_folder/venv_name/images"


# Perform object detection on every image in specified folder
results = model(
    source=image_folder,
    show=False,  # Set to True if you want to display the image
    cnotallow=0.60,  # Confidence threshold
    save=True,  # Save results
    save_txt=True,  # Save results as text files
    save_cnotallow=True,  # Save confidence scores
    line_width=3,  # Adjust line width for bounding boxes
    save_crop=True  # Save cropped detections as image files
    )

由于每张图像的目标检测结果都保存到单独的文本文件中，你可以使用pandas读取所有这些文件以获取检测到的对象总数。结果文件还包含每个检测的类别、x_center、y_center、宽度和高度，这些数据可能对进一步分析有用。

复制

import os
import pandas as pd


# Path to folder containing result text files
folder_path = "/home/username/venv_folder/venv_name/runs/detect/predict/labels"


# Empty list to store data from each file
dataframes = []


# Define column names
column_names = ["class", "x_center", "y_center", "width", "height", "confidence"]


# Iterate over all result files in the folder
for filename in os.listdir(folder_path):
    file_path = os.path.join(folder_path, filename)
    df = pd.read_csv(file_path, delimiter=' ', header=None, names=column_names,)
    df['filename'] = filename  # Add column to record filename of each result
    dataframes.append(df) # Append results to dataframes list


# Combine all into a single DataFrame
combined_df = pd.concat(dataframes, ignore_index=True)


# Check the top five rows of the DataFrame
display(combined_df.head(5))


# Count number of rows in DataFrame
num_rows = combined_df.shape[0]
print(f"Number of detected stamps: {num_rows}")

自定义训练的 YOLOv8 模型进行邮票整理