使用YOLO的TXT目標(biāo)檢測(cè)標(biāo)簽的滑窗切割：批量處理圖像和標(biāo)簽的實(shí)用工具

背景

在計(jì)算機(jī)視覺(jué)領(lǐng)域，目標(biāo)檢測(cè)（Object Detection）是一個(gè)非常重要的任務(wù)。隨著 YOLO（You Only Look Once）系列模型的普及，目標(biāo)檢測(cè)模型已經(jīng)被廣泛應(yīng)用于各種實(shí)際場(chǎng)景中。對(duì)于目標(biāo)檢測(cè)任務(wù)，訓(xùn)練模型所需的標(biāo)注數(shù)據(jù)至關(guān)重要。

當(dāng)我們處理大規(guī)模圖像數(shù)據(jù)集時(shí)，尤其是在圖像的尺寸遠(yuǎn)大于模型輸入尺寸時(shí)，往往需要使用 滑窗切割（Sliding Window）技術(shù)，將大圖像分割成多個(gè)小塊進(jìn)行處理。這一過(guò)程不僅可以減小每次訓(xùn)練所需的計(jì)算資源，還能增強(qiáng)模型的魯棒性。

本博客將介紹如何使用 YOLO 的TXT目標(biāo)檢測(cè)標(biāo)簽格式 對(duì)大圖像進(jìn)行滑窗切割，并確保標(biāo)簽的正確性。我們將逐步闡述該代碼的工作原理、使用方法及其在目標(biāo)檢測(cè)中的實(shí)際意義。

1. 代碼概述

該代碼實(shí)現(xiàn)了對(duì)大圖像及其對(duì)應(yīng)標(biāo)簽的 滑窗切割，并確保切割后的標(biāo)簽正確地被裁剪并保存。它通過(guò)對(duì)圖像和標(biāo)簽的逐塊切割，將大圖像分割成多個(gè)較小的圖像塊，同時(shí)調(diào)整標(biāo)簽的位置和大小，以符合新的圖像尺寸。

主要步驟如下：

1. 加載圖像和標(biāo)簽：讀取圖片和標(biāo)簽文件，確保標(biāo)簽與圖像對(duì)應(yīng)。
2. 滑窗切割：以給定的窗口大小和步長(zhǎng)，對(duì)圖像進(jìn)行滑窗切割。
3. 裁剪標(biāo)簽：對(duì)于每個(gè)切割窗口，檢查標(biāo)簽是否位于窗口內(nèi)，如果位于窗口內(nèi)，調(diào)整標(biāo)簽坐標(biāo)，并確保標(biāo)簽歸一化。
4. 保存切割后的圖像和標(biāo)簽：將切割后的圖像和標(biāo)簽保存到新的文件夾中。

2. 滑窗切割算法原理

滑窗切割是計(jì)算機(jī)視覺(jué)中常用的技術(shù)，通常用于：

• 大圖像分塊：當(dāng)圖像尺寸過(guò)大時(shí)，模型輸入尺寸無(wú)法處理整個(gè)圖像，可以將其切割成小塊進(jìn)行逐塊處理。
• 多尺度檢測(cè)：不同尺度的物體需要不同大小的窗口來(lái)檢測(cè)。通過(guò)滑窗切割，能夠在多個(gè)尺度上執(zhí)行目標(biāo)檢測(cè)任務(wù)。

滑窗切割步驟：

1. 指定窗口大小和步長(zhǎng)：窗口大小和步長(zhǎng)決定了滑窗的密集程度。步長(zhǎng)越小，生成的窗口越多，計(jì)算量越大。窗口大小決定了每個(gè)塊的輸入尺寸。

2. 標(biāo)簽裁剪：標(biāo)簽的裁剪是根據(jù)目標(biāo)與滑窗的交集來(lái)進(jìn)行的。每個(gè)標(biāo)簽會(huì)被裁剪到窗口內(nèi)，并且坐標(biāo)會(huì)被重新歸一化到窗口的尺寸。

示例：

• 窗口大小：640x640像素。
• 橫向步長(zhǎng)：301像素。
• 縱向步長(zhǎng)：180像素。

對(duì)于每個(gè)標(biāo)簽，代碼會(huì)檢查它是否位于當(dāng)前滑窗內(nèi)，如果是，標(biāo)簽的位置和尺寸會(huì)被重新計(jì)算并保存。

3. 代碼實(shí)現(xiàn)

1. 加載標(biāo)簽

def load_labels(label_file):"""加載YOLO的標(biāo)簽文件"""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

這段代碼用于讀取每個(gè)標(biāo)簽文件，并將其轉(zhuǎn)換為包含類別和坐標(biāo)的格式，方便后續(xù)處理。

2. 切割標(biāo)簽

def save_cut_labels(window_x, window_y, window_size, img_width, img_height, labels):"""根據(jù)滑窗切割標(biāo)簽，并確保標(biāo)簽正確裁剪"""new_labels = []for cls, x_center, y_center, w, h in labels:# 將歸一化坐標(biāo)轉(zhuǎn)換為像素坐標(biāo)x_center_px = x_center * img_widthy_center_px = y_center * img_heightw_px = w * img_widthh_px = h * img_height# 計(jì)算標(biāo)簽與當(dāng)前窗口的交集區(qū)域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)# 如果標(biāo)簽和窗口相交if intersection_x1 < intersection_x2 and intersection_y1 < intersection_y2:# 計(jì)算交集區(qū)域的寬高和中心坐標(biāo)intersection_w = intersection_x2 - intersection_x1intersection_h = intersection_y2 - intersection_y1intersection_x_center = (intersection_x1 + intersection_x2) / 2intersection_y_center = (intersection_y1 + intersection_y2) / 2# 將交集區(qū)域的坐標(biāo)歸一化normalized_x_center = (intersection_x_center - window_x) / window_sizenormalized_y_center = (intersection_y_center - window_y) / window_sizenormalized_w = intersection_w / window_sizenormalized_h = intersection_h / window_size# 生成新的標(biāo)簽new_labels.append(f"{cls} {normalized_x_center} {normalized_y_center} {normalized_w} {normalized_h}")return new_labels

該函數(shù)根據(jù)當(dāng)前窗口的位置，裁剪標(biāo)簽，并將裁剪后的標(biāo)簽歸一化到當(dāng)前窗口大小。

3. 主函數(shù)

def main():image_folder = 'images'  # 輸入圖片文件夾label_folder = 'labels'  # 輸入標(biāo)簽文件夾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  # 橫向步長(zhǎng)step_y = 180  # 縱向步長(zhǎng)# 遍歷所有圖片和標(biāo)簽文件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# 讀取對(duì)應(yīng)的標(biāo)簽label_path = os.path.join(label_folder, label_file)labels = load_labels(label_path)# 計(jì)算橫向和縱向可以切割的窗口數(shù)量num_windows_x = (img_width - window_size) // step_x + 1num_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_xwindow_y = j * step_y# 獲取當(dāng)前窗口內(nèi)的標(biāo)簽windowed_labels = save_cut_labels(window_x, window_y, window_size, img_width, img_height, labels)if windowed_labels:  # 如果窗口內(nèi)有標(biāo)簽# 保存切割后的圖片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)# 保存切割后的標(biāo)簽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. 準(zhǔn)備工作：

   • 將你的圖片和標(biāo)簽放在 images/ 和 labels/ 文件夾中。
   • 確保標(biāo)簽格式為 YOLOv5 格式，即每行包含 class_id x_center y_center width height（所有值均為歸一化形式）。

2. 運(yùn)行腳本：

   • 運(yùn)行上述代碼，程序?qū)⒆詣?dòng)讀取圖片和標(biāo)簽，進(jìn)行滑窗切割，并將每個(gè)切割后的圖像和標(biāo)簽保存到新的文件夾中。

3. 輸出結(jié)果：

   • 切割后的圖像會(huì)保存在 output_images/ 文件夾中。
   • 切割后的標(biāo)簽會(huì)保存在 output_labels/ 文件夾中，標(biāo)簽內(nèi)容與原標(biāo)簽一致，只是經(jīng)過(guò)裁

4. 完整代碼

import os
import cv2def load_labels(label_path):"""加載YOLOv5標(biāo)簽文件"""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 labelsdef save_cut_labels(window_x, window_y, window_size, img_width, img_height, labels):"""根據(jù)滑窗切割標(biāo)簽，并確保標(biāo)簽正確裁剪"""new_labels = []for cls, x_center, y_center, w, h in labels:# 將歸一化坐標(biāo)轉(zhuǎn)換為像素坐標(biāo)x_center_px = x_center * img_widthy_center_px = y_center * img_heightw_px = w * img_widthh_px = h * img_height# 計(jì)算標(biāo)簽與當(dāng)前窗口的交集區(qū)域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)# 如果標(biāo)簽和窗口相交if intersection_x1 < intersection_x2 and intersection_y1 < intersection_y2:# 計(jì)算交集區(qū)域的寬高和中心坐標(biāo)intersection_w = intersection_x2 - intersection_x1intersection_h = intersection_y2 - intersection_y1intersection_x_center = (intersection_x1 + intersection_x2) / 2intersection_y_center = (intersection_y1 + intersection_y2) / 2# 將交集區(qū)域的坐標(biāo)歸一化normalized_x_center = (intersection_x_center - window_x) / window_sizenormalized_y_center = (intersection_y_center - window_y) / window_sizenormalized_w = intersection_w / window_sizenormalized_h = intersection_h / window_size# 生成新的標(biāo)簽new_labels.append(f"{cls} {normalized_x_center} {normalized_y_center} {normalized_w} {normalized_h}")# 如果沒(méi)有標(biāo)簽，返回空列表return new_labelsdef main():image_folder = 'images'  # 輸入圖片文件夾label_folder = 'labels'  # 輸入標(biāo)簽文件夾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  # 橫向步長(zhǎng)step_y = 180  # 縱向步長(zhǎng)# 遍歷所有圖片和標(biāo)簽文件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# 讀取對(duì)應(yīng)的標(biāo)簽label_path = os.path.join(label_folder, label_file)labels = load_labels(label_path)# 計(jì)算橫向和縱向可以切割的窗口數(shù)量num_windows_x = (img_width - window_size) // step_x + 1num_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_xwindow_y = j * step_y# 獲取當(dāng)前窗口內(nèi)的標(biāo)簽windowed_labels = save_cut_labels(window_x, window_y, window_size, img_width, img_height, labels)# 如果標(biāo)簽列表為空，說(shuō)明此窗口沒(méi)有標(biāo)簽，跳過(guò)該窗口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)# 保存切割后的標(biāo)簽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()