【目标检测】根据YOLO格式标签提取和保存标注框目标（附完整代码）

关于作者

作者：小白熊

作者简介：精通python、matlab、c#语言，擅长机器学习，深度学习，机器视觉，目标检测，图像分类，姿态识别，语义分割，路径规划，智能优化算法，数据分析，各类创新融合等等。

联系邮箱：xbx3144@163.com

科研辅导、知识付费答疑、个性化定制以及其他合作需求请联系作者~

1 前言

在这篇博文中，我们将讨论如何通过Python实现从YOLO标注文件中提取目标检测框，并将检测出的目标区域从图片中截取并保存为单独的图像文件。YOLO的标注格式包含的标注信息包括目标类别和其归一化后的位置与大小，以.txt格式存储，每行表示一个目标，包含以下字段：

class_id：目标类别编号
x_center：目标框中心点的横坐标（相对于图像宽度归一化）
y_center：目标框中心点的纵坐标（相对于图像高度归一化）
width：目标框的宽度（相对于图像宽度归一化）
height：目标框的高度（相对于图像高度归一化）

本文代码的目标是读取YOLO标注文件，将标注的目标从原始图片中截取出来，并按类别分类保存。

2 代码实现

2.1 依赖库

import os
import cv2
import warningswarnings.filterwarnings("ignore")

os用于处理文件和路径，OpenCV帮助我们处理图像和执行图像相关的操作。同时，使用warnings.filterwarnings("ignore")来忽略警告信息，以保持输出的清晰。

2.2 设置路径

img_folder = "./images" 
txt_folder = "./labels"

这里定义了两个文件夹路径：

img_folder：存储待处理的图片。
txt_folder：存储对应图片的YOLO标注文件。

2.3 遍历文件

for filename in os.listdir(img_folder):if filename.endswith(".jpg") or filename.endswith(".jpeg") or filename.endswith(".png"):image_file_path = os.path.join(img_folder, filename)txt_filename = os.path.splitext(filename)[0] + ".txt"label_file_path = os.path.join(txt_folder, txt_filename)

我们通过os.listdir()函数遍历图片文件夹中的文件，并确保只处理图像文件。对应的标注文件名与图像文件名相同，扩展名为.txt。

2.4 读取标签

with open(label_file_path, 'r') as f:lines = f.readlines()
image = cv2.imread(image_file_path)
i = 0
for line in lines:i = i + 1class_id, x_center, y_center, width, height = map(float, line.strip().split())

通过open()函数读取与图像文件名相同的标注文件内容，lines变量存储了文件中的每一行标注信息。同时使用cv2.imread()读取图片文件，并将其保存为image变量。接着开始逐行处理标注文件中的每一行数据。每行数据用空格分隔，表示类别编号、目标框的中心坐标、宽度和高度，并将这些数据转换为浮点型数值。

2.5 坐标转换

left = int((x_center - width / 2) * image.shape[1])
top = int((y_center - height / 2) * image.shape[0])
right = int((x_center + width / 2) * image.shape[1])
bottom = int((y_center + height / 2) * image.shape[0])

YOLO坐标是相对于图像宽度和高度的归一化值，因此我们需要将其转换为实际的图像像素坐标：

left和right分别是目标框的左边界和右边界坐标。
top和bottom分别是目标框的上边界和下边界坐标。

这些计算通过使用图像的宽度image.shape[1]和高度image.shape[0]来完成。

2.6 保存图片

class_id = int(class_id)if class_id == 0:folder = 'label 1'
elif class_id == 1:folder = 'label 2'
elif class_id == 2:folder = 'label 3'object_image = image[top:bottom, left:right]
object_file_path = image_file_path.split("\\")
object_file_path = f"Target box extraction/{folder}/" + os.path.splitext(object_file_path[1])[0] + f"_{int(class_id)}_{i}.jpg"
cv2.imwrite(object_file_path, object_image)

通过class_id来判断目标所属的类别，根据不同的类别，保存到不同的文件夹中。例如，类别编号为0的目标保存到label 1文件夹，类别编号为1的保存到label 2文件夹，依此类推。需要根据自己项目的实际情况进行修改。

然后使用image[top:bottom, left:right]从原始图像中截取出目标框区域。object_file_path用于构建保存新图像的路径，该路径包含目标的类别和编号。最后，使用cv2.imwrite()函数将截取出的目标保存为新图像。

3 完整代码

import os
import cv2
import warningswarnings.filterwarnings("ignore")# 指定图片文件夹路径和txt文件夹路径
img_folder = "./images"
txt_folder = "./labels"# 遍历图片文件夹中的所有文件
for filename in os.listdir(img_folder):# 确保文件是一个图像文件if filename.endswith(".jpg") or filename.endswith(".jpeg") or filename.endswith(".png"):# 构造图像文件路径image_file_path = os.path.join(img_folder, filename)# 构造对应的txt文件路径txt_filename = os.path.splitext(filename)[0] + ".txt"label_file_path = os.path.join(txt_folder, txt_filename)# 读取标注文件和原始图片with open(label_file_path, 'r') as f:lines = f.readlines()image = cv2.imread(image_file_path)# 循环处理每个标注框i = 0for line in lines:i = i + 1class_id, x_center, y_center, width, height = map(float, line.strip().split())# 将YOLO格式的坐标转换为常规坐标left = int((x_center - width / 2) * image.shape[1])top = int((y_center - height / 2) * image.shape[0])right = int((x_center + width / 2) * image.shape[1])bottom = int((y_center + height / 2) * image.shape[0])class_id = int(class_id)if class_id == 0:folder = 'label 1'elif class_id == 1:folder = 'label 2'elif class_id == 2:folder = 'label 3'# 截取标注框内的内容并保存为新图片object_image = image[top:bottom, left:right]object_file_path = image_file_path.split("\\")object_file_path = f"Target box extraction/{folder}/" + os.path.splitext(object_file_path[1])[0] + f"_{int(class_id)}_{i}.jpg"cv2.imwrite(object_file_path, object_image)