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

qaz

还记得人们过去常常寄信和明信片的日子吗？明信片上通常描绘了你所访问的城市或国家的一些风景如画的景色，但对许多人来说，邮票才是主要的吸引力。它们被视为微型艺术品，被认为是非常值得收藏的。
由于每张明信片、信封、包裹和邮包都需要邮票，许多人最终拥有了大量的邮票收藏。这变成了一个计算机视觉项目：我没有一张一张地拍摄每张邮票的照片，而是拍摄了邮票相册的每一页，并使用目标检测来找到并存储图片中所有邮票的图像。

目标检测

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

使用自定义训练的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 datasettest: test/images # relative path to test imagestrain: train/images # relative path to training imagesval: val/images # relative path to validation images# classesnames:  0: postage-stamp

• 1.
  
• 2.
  
• 3.
  
• 4.
  
• 5.
  
• 6.
  
• 7.
  
• 8.
  

使用自定义数据集训练YOLOv8所需的Python代码如下所示。建议使用GPU，并在训练前检查其可用性。
import torch, ultralytics# Check library versionsprint("PyTorch version:", torch.__version__)print("Ultralytics version:", ultralytics.__version__)# Check if GPU is availableif torch.cuda.is_available():    print(f"GPU is available: {torch.cuda.get_device_name(0)}")else:    print("GPU is not available")

• 1.
  
• 2.
  
• 3.
  
• 4.
  
• 5.
  
• 6.
  
• 7.
  
• 8.
  
• 9.
  
• 10.
  
• 11.
  
• 12.
  
• 13.
  

在训练之前，加载带有预训练权重的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 modelmodel.train(data="config.yaml", epochs=100, patience=10)

• 1.
  
• 2.
  
• 3.
  
• 4.
  
• 5.
  
• 6.
  
• 7.
  
• 8.
  
• 9.
  
• 10.
  
• 11.
  
• 12.
  
• 13.
  
• 14.
  
• 15.
  
• 16.
  
• 17.
  

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

YOLOv8训练指标
要使用自定义训练的模型从新图像中检测对象，请从runs/detect/train/weights文件夹加载它。你可以选择best.pt或last.pt，前者是得分最高的模型，后者是最后一个epoch训练的模型。
# Load the trained YOLO modelmodel = YOLO("/home/username/venv_folder/venv_name/runs/detect/train/weights/best.pt")# Specify folder containing images for object detectionimage_folder = "/home/username/venv_folder/venv_name/images"# Perform object detection on every image in specified folderresults = 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    )

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

由于每张图像的目标检测结果都保存到单独的文本文件中，你可以使用pandas读取所有这些文件以获取检测到的对象总数。结果文件还包含每个检测的类别、x_center、y_center、宽度和高度，这些数据可能对进一步分析有用。
import osimport pandas as pd# Path to folder containing result text filesfolder_path = "/home/username/venv_folder/venv_name/runs/detect/predict/labels"# Empty list to store data from each filedataframes = []# Define column namescolumn_names = ["class", "x_center", "y_center", "width", "height", "confidence"]# Iterate over all result files in the folderfor 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 DataFramecombined_df = pd.concat(dataframes, ignore_index=True)# Check the top five rows of the DataFramedisplay(combined_df.head(5))# Count number of rows in DataFramenum_rows = combined_df.shape[0]print(f"Number of detected stamps: {num_rows}")

• 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.
  

检测结果的DataFrame（前五行）和检测到的对象计数
人脸检测

现在，收藏中每张邮票的图像都已保存（在我的情况下总共945张图像），如何搜索其中包含人脸的邮票呢？没问题，使用人脸数据集训练新模型并再次运行目标检测。

 

 

使用人脸数据集训练的YOLOv8n检测邮票中的人脸
一些看起来更像卡通的人脸没有被检测到，因为训练数据由实际人脸的图片组成。需要进一步调整以获得更好的性能，但希望我能够展示使用自定义数据集训练YOLO检测任何对象是多么容易。

参考资料

• [1] https://universe.roboflow.com/jackwildgooglecom/detect-postage-stamp
  
• [2] https://docs.ultralytics.com/models/yolov8/#supported-tasks-and-modes