使用YOLO系列txt目标检测标签的滑窗切割:批量处理图像和标签的实用工具

使用YOLO系列txt目标检测标签的滑窗切割:批量处理图像和标签的实用工具

使用YOLO的TXT目标检测标签的滑窗切割:批量处理图像和标签的实用工具

背景

在计算机视觉领域,目标检测 (Object Detection)是一个非常重要的任务。随着 YOLO(You Only Look Once)系列模型的普及,目标检测模型已经被广泛应用于各种实际场景中。对于目标检测任务,训练模型所需的标注数据至关重要。

当我们处理大规模图像数据集时,尤其是在图像的尺寸远大于模型输入尺寸时,往往需要使用 滑窗切割(Sliding Window)技术,将大图像分割成多个小块进行处理。这一过程不仅可以减小每次训练所需的计算资源,还能增强模型的鲁棒性。

本博客将介绍如何使用 YOLO 的TXT目标检测标签格式 对大图像进行滑窗切割,并确保标签的正确性。我们将逐步阐述该代码的工作原理、使用方法及其在目标检测中的实际意义。

1. 代码概述

该代码实现了对大图像及其对应标签的 滑窗切割,并确保切割后的标签正确地被裁剪并保存。它通过对图像和标签的逐块切割,将大图像分割成多个较小的图像块,同时调整标签的位置和大小,以符合新的图像尺寸。

主要步骤如下:

  1. 加载图像和标签:读取图片和标签文件,确保标签与图像对应。
  2. 滑窗切割:以给定的窗口大小和步长,对图像进行滑窗切割。
  3. 裁剪标签:对于每个切割窗口,检查标签是否位于窗口内,如果位于窗口内,调整标签坐标,并确保标签归一化。
  4. 保存切割后的图像和标签:将切割后的图像和标签保存到新的文件夹中。

2. 滑窗切割算法原理

滑窗切割是计算机视觉中常用的技术,通常用于:

  • 大图像分块:当图像尺寸过大时,模型输入尺寸无法处理整个图像,可以将其切割成小块进行逐块处理。
  • 多尺度检测:不同尺度的物体需要不同大小的窗口来检测。通过滑窗切割,能够在多个尺度上执行目标检测任务。
滑窗切割步骤:
  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. 如何使用该工具

  1. 准备工作

    • 将你的图片和标签放在 images/labels/ 文件夹中。
    • 确保标签格式为 YOLOv5 格式,即每行包含 class_id x_center y_center width height(所有值均为归一化形式)。
  2. 运行脚本

    • 运行上述代码,程序将自动读取图片和标签,进行滑窗切割,并将每个切割后的图像和标签保存到新的文件夹中。
  3. 输出结果

    • 切割后的图像会保存在 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()
相关推荐
Donvink15 分钟前
Transformers在计算机视觉领域中的应用【第3篇:Swin Transformer——多层次的Vision Transformer】
人工智能·深度学习·目标检测·计算机视觉·transformer
龙的爹233320 分钟前
2024论文翻译 | Multi-Review Fusion-in-Context
人工智能·深度学习·自然语言处理·prompt
一尘之中1 小时前
基于Transformer的编码器-解码器图像描述模型在AMD GPU上的应用
人工智能·深度学习·transformer
Seeklike2 小时前
12.04 深度学习-用CNN做图像分类+训练可视化
深度学习·分类·cnn
亦枫Leonlew7 小时前
三维测量与建模笔记 - 5.3 光束法平差(Bundle Adjustment)
笔记·计算机视觉·三维重建·光束法平差
爱研究的小牛9 小时前
Runway 技术浅析(七):视频技术中的运动跟踪
人工智能·深度学习·计算机视觉·目标跟踪·aigc
DieYoung_Alive9 小时前
搭建深度学习框架+nn.Module
人工智能·深度学习·yolo
GOTXX9 小时前
修改训练策略,无损提升性能
人工智能·计算机视觉·目标跟踪
B站计算机毕业设计超人12 小时前
计算机毕业设计Python异常流量检测 流量分类 流量分析 网络流量分析与可视化系统 网络安全 信息安全 机器学习 深度学习
大数据·人工智能·python·深度学习·web安全·机器学习·课程设计
湘人-汤义12 小时前
transformers实现一个检索机器人(一)
人工智能·深度学习·自然语言处理·机器人