使用YOLO系列txt目标检测标签的滑窗切割:批量处理图像和标签的实用工具
使用YOLO的TXT目标检测标签的滑窗切割:批量处理图像和标签的实用工具
背景
在计算机视觉领域,目标检测 (Object Detection)是一个非常重要的任务。随着 YOLO(You Only Look Once)系列模型的普及,目标检测模型已经被广泛应用于各种实际场景中。对于目标检测任务,训练模型所需的标注数据至关重要。
当我们处理大规模图像数据集时,尤其是在图像的尺寸远大于模型输入尺寸时,往往需要使用 滑窗切割(Sliding Window)技术,将大图像分割成多个小块进行处理。这一过程不仅可以减小每次训练所需的计算资源,还能增强模型的鲁棒性。
本博客将介绍如何使用 YOLO 的TXT目标检测标签格式 对大图像进行滑窗切割,并确保标签的正确性。我们将逐步阐述该代码的工作原理、使用方法及其在目标检测中的实际意义。
1. 代码概述
该代码实现了对大图像及其对应标签的 滑窗切割,并确保切割后的标签正确地被裁剪并保存。它通过对图像和标签的逐块切割,将大图像分割成多个较小的图像块,同时调整标签的位置和大小,以符合新的图像尺寸。
主要步骤如下:
- 加载图像和标签:读取图片和标签文件,确保标签与图像对应。
- 滑窗切割:以给定的窗口大小和步长,对图像进行滑窗切割。
- 裁剪标签:对于每个切割窗口,检查标签是否位于窗口内,如果位于窗口内,调整标签坐标,并确保标签归一化。
- 保存切割后的图像和标签:将切割后的图像和标签保存到新的文件夹中。
2. 滑窗切割算法原理
滑窗切割是计算机视觉中常用的技术,通常用于:
- 大图像分块:当图像尺寸过大时,模型输入尺寸无法处理整个图像,可以将其切割成小块进行逐块处理。
- 多尺度检测:不同尺度的物体需要不同大小的窗口来检测。通过滑窗切割,能够在多个尺度上执行目标检测任务。
滑窗切割步骤:
-
指定窗口大小和步长:窗口大小和步长决定了滑窗的密集程度。步长越小,生成的窗口越多,计算量越大。窗口大小决定了每个块的输入尺寸。
-
标签裁剪:标签的裁剪是根据目标与滑窗的交集来进行的。每个标签会被裁剪到窗口内,并且坐标会被重新归一化到窗口的尺寸。
示例:
- 窗口大小:640x640像素。
- 横向步长:301像素。
- 纵向步长:180像素。
对于每个标签,代码会检查它是否位于当前滑窗内,如果是,标签的位置和尺寸会被重新计算并保存。
3. 代码实现
1. 加载标签
python
def load_labels(label_file):
"""加载YOLO的标签文件"""
labels = []
with open(label_file, 'r') as f:
for line in f:
parts = line.strip().split()
cls = int(parts[0]) # 类别
x_center, y_center, w, h = map(float, parts[1:])
labels.append((cls, x_center, y_center, w, h))
return labels
这段代码用于读取每个标签文件,并将其转换为包含类别和坐标的格式,方便后续处理。
2. 切割标签
python
def save_cut_labels(window_x, window_y, window_size, img_width, img_height, labels):
"""根据滑窗切割标签,并确保标签正确裁剪"""
new_labels = []
for cls, x_center, y_center, w, h in labels:
# 将归一化坐标转换为像素坐标
x_center_px = x_center * img_width
y_center_px = y_center * img_height
w_px = w * img_width
h_px = h * img_height
# 计算标签与当前窗口的交集区域
intersection_x1 = max(x_center_px - w_px / 2, window_x)
intersection_y1 = max(y_center_px - h_px / 2, window_y)
intersection_x2 = min(x_center_px + w_px / 2, window_x + window_size)
intersection_y2 = min(y_center_px + h_px / 2, window_y + window_size)
# 如果标签和窗口相交
if intersection_x1 < intersection_x2 and intersection_y1 < intersection_y2:
# 计算交集区域的宽高和中心坐标
intersection_w = intersection_x2 - intersection_x1
intersection_h = intersection_y2 - intersection_y1
intersection_x_center = (intersection_x1 + intersection_x2) / 2
intersection_y_center = (intersection_y1 + intersection_y2) / 2
# 将交集区域的坐标归一化
normalized_x_center = (intersection_x_center - window_x) / window_size
normalized_y_center = (intersection_y_center - window_y) / window_size
normalized_w = intersection_w / window_size
normalized_h = intersection_h / window_size
# 生成新的标签
new_labels.append(f"{cls} {normalized_x_center} {normalized_y_center} {normalized_w} {normalized_h}")
return new_labels
该函数根据当前窗口的位置,裁剪标签,并将裁剪后的标签归一化到当前窗口大小。
3. 主函数
python
def main():
image_folder = 'images' # 输入图片文件夹
label_folder = 'labels' # 输入标签文件夹
output_image_folder = 'output_images'
output_label_folder = 'output_labels'
if not os.path.exists(output_image_folder):
os.makedirs(output_image_folder)
if not os.path.exists(output_label_folder):
os.makedirs(output_label_folder)
image_files = sorted(os.listdir(image_folder))
label_files = sorted(os.listdir(label_folder))
window_size = 640 # 滑窗大小
step_x = 301 # 横向步长
step_y = 180 # 纵向步长
# 遍历所有图片和标签文件
for image_file, label_file in zip(image_files, label_files):
# 读取图片
image_path = os.path.join(image_folder, image_file)
image = cv2.imread(image_path)
img_height, img_width, _ = image.shape
# 读取对应的标签
label_path = os.path.join(label_folder, label_file)
labels = load_labels(label_path)
# 计算横向和纵向可以切割的窗口数量
num_windows_x = (img_width - window_size) // step_x + 1
num_windows_y = (img_height - window_size) // step_y + 1
# 遍历所有切割窗口
for i in range(num_windows_x):
for j in range(num_windows_y):
window_x = i * step_x
window_y = j * step_y
# 获取当前窗口内的标签
windowed_labels = save_cut_labels(window_x, window_y, window_size, img_width, img_height, labels)
if windowed_labels: # 如果窗口内有标签
# 保存切割后的图片
windowed_image = image[window_y:window_y + window_size, window_x:window_x + window_size]
output_image_path = os.path.join(output_image_folder, f"{os.path.splitext(image_file)[0]}_window_{i}_{j}.jpg")
cv2.imwrite(output_image_path, windowed_image)
# 保存切割后的标签
output_label_path = os.path.join(output_label_folder, f"{os.path.splitext(label_file)[0]}_window_{i}_{j}.txt")
with open(output_label_path, 'w') as f:
for label in windowed_labels:
f.write(label + '\n')
4. 如何使用该工具
-
准备工作:
- 将你的图片和标签放在
images/
和labels/
文件夹中。 - 确保标签格式为 YOLOv5 格式,即每行包含
class_id x_center y_center width height
(所有值均为归一化形式)。
- 将你的图片和标签放在
-
运行脚本:
- 运行上述代码,程序将自动读取图片和标签,进行滑窗切割,并将每个切割后的图像和标签保存到新的文件夹中。
-
输出结果:
- 切割后的图像会保存在
output_images/
文件夹中。 - 切割后的标签会保存在
output_labels/
文件夹中,标签内容与原标签一致,只是经过裁
- 切割后的图像会保存在
4. 完整代码
python
import os
import cv2
def load_labels(label_path):
"""加载YOLOv5标签文件"""
labels = []
with open(label_path, 'r') as f:
for line in f.readlines():
parts = line.strip().split()
cls = int(parts[0]) # 类别
x_center = float(parts[1]) # x中心
y_center = float(parts[2]) # y中心
w = float(parts[3]) # 宽度
h = float(parts[4]) # 高度
labels.append([cls, x_center, y_center, w, h])
return labels
def save_cut_labels(window_x, window_y, window_size, img_width, img_height, labels):
"""根据滑窗切割标签,并确保标签正确裁剪"""
new_labels = []
for cls, x_center, y_center, w, h in labels:
# 将归一化坐标转换为像素坐标
x_center_px = x_center * img_width
y_center_px = y_center * img_height
w_px = w * img_width
h_px = h * img_height
# 计算标签与当前窗口的交集区域
intersection_x1 = max(x_center_px - w_px / 2, window_x)
intersection_y1 = max(y_center_px - h_px / 2, window_y)
intersection_x2 = min(x_center_px + w_px / 2, window_x + window_size)
intersection_y2 = min(y_center_px + h_px / 2, window_y + window_size)
# 如果标签和窗口相交
if intersection_x1 < intersection_x2 and intersection_y1 < intersection_y2:
# 计算交集区域的宽高和中心坐标
intersection_w = intersection_x2 - intersection_x1
intersection_h = intersection_y2 - intersection_y1
intersection_x_center = (intersection_x1 + intersection_x2) / 2
intersection_y_center = (intersection_y1 + intersection_y2) / 2
# 将交集区域的坐标归一化
normalized_x_center = (intersection_x_center - window_x) / window_size
normalized_y_center = (intersection_y_center - window_y) / window_size
normalized_w = intersection_w / window_size
normalized_h = intersection_h / window_size
# 生成新的标签
new_labels.append(f"{cls} {normalized_x_center} {normalized_y_center} {normalized_w} {normalized_h}")
# 如果没有标签,返回空列表
return new_labels
def main():
image_folder = 'images' # 输入图片文件夹
label_folder = 'labels' # 输入标签文件夹
output_image_folder = 'output_images'
output_label_folder = 'output_labels'
if not os.path.exists(output_image_folder):
os.makedirs(output_image_folder)
if not os.path.exists(output_label_folder):
os.makedirs(output_label_folder)
image_files = sorted(os.listdir(image_folder))
label_files = sorted(os.listdir(label_folder))
window_size = 640 # 滑窗大小
step_x = 301 # 横向步长
step_y = 180 # 纵向步长
# 遍历所有图片和标签文件
for image_file, label_file in zip(image_files, label_files):
# 读取图片
image_path = os.path.join(image_folder, image_file)
image = cv2.imread(image_path)
img_height, img_width, _ = image.shape
# 读取对应的标签
label_path = os.path.join(label_folder, label_file)
labels = load_labels(label_path)
# 计算横向和纵向可以切割的窗口数量
num_windows_x = (img_width - window_size) // step_x + 1
num_windows_y = (img_height - window_size) // step_y + 1
# 遍历所有切割窗口
for i in range(num_windows_x):
for j in range(num_windows_y):
window_x = i * step_x
window_y = j * step_y
# 获取当前窗口内的标签
windowed_labels = save_cut_labels(window_x, window_y, window_size, img_width, img_height, labels)
# 如果标签列表为空,说明此窗口没有标签,跳过该窗口
if not windowed_labels:
continue
# 保存切割后的图片
windowed_image = image[window_y:window_y + window_size, window_x:window_x + window_size]
output_image_path = os.path.join(output_image_folder, f"{os.path.splitext(image_file)[0]}_window_{i}_{j}.jpg")
cv2.imwrite(output_image_path, windowed_image)
# 保存切割后的标签
output_label_path = os.path.join(output_label_folder, f"{os.path.splitext(label_file)[0]}_window_{i}_{j}.txt")
with open(output_label_path, 'w') as f:
for label in windowed_labels:
f.write(label + '\n')
if __name__ == "__main__":
main()