作者:CSDN @ _養樂多_
本文將介紹如何將目標檢測中常用的VOC格式數據集轉換為YOLO數據集,并進行數據集比例劃分,從而方便的進行YOLO目標檢測。
如果不想分兩步,可以直接看第三節代碼。
文章目錄
- 一、將VOC格式數據集轉換為YOLO格式數據集
- 二、YOLO格式數據集劃分(訓練、驗證、測試)
- 2.1 版本1
- 2.2 版本2
- 2.3 版本3
- 三、一步到位
一、將VOC格式數據集轉換為YOLO格式數據集
執行以下腳本將VOC格式數據集轉換為YOLO格式數據集。
但是需要注意的是:
- 轉換之后的數據集只有Images和labels兩個文件。還需要執行第二節中的腳本進行數據集劃分,將總的數據集劃分為訓練、驗證、測試數據集;
- 使用的話,需要修改 class_mapping 中類別名和對應標簽,還有VOC數據集路徑、YOLO數據集路徑。
import os
import shutil
import xml.etree.ElementTree as ET# VOC格式數據集路徑
voc_data_path = 'E:\\DataSet\\helmet-VOC'
voc_annotations_path = os.path.join(voc_data_path, 'Annotations')
voc_images_path = os.path.join(voc_data_path, 'JPEGImages')# YOLO格式數據集保存路徑
yolo_data_path = 'E:\\DataSet\\helmet-YOLO'
yolo_images_path = os.path.join(yolo_data_path, 'images')
yolo_labels_path = os.path.join(yolo_data_path, 'labels')# 創建YOLO格式數據集目錄
os.makedirs(yolo_images_path, exist_ok=True)
os.makedirs(yolo_labels_path, exist_ok=True)# 類別映射 (可以根據自己的數據集進行調整)
class_mapping = {'head': 0,'helmet': 1,'person': 2,# 添加更多類別...
}def convert_voc_to_yolo(voc_annotation_file, yolo_label_file):tree = ET.parse(voc_annotation_file)root = tree.getroot()size = root.find('size')width = float(size.find('width').text)height = float(size.find('height').text)with open(yolo_label_file, 'w') as f:for obj in root.findall('object'):cls = obj.find('name').textif cls not in class_mapping:continuecls_id = class_mapping[cls]xmlbox = obj.find('bndbox')xmin = float(xmlbox.find('xmin').text)ymin = float(xmlbox.find('ymin').text)xmax = float(xmlbox.find('xmax').text)ymax = float(xmlbox.find('ymax').text)x_center = (xmin + xmax) / 2.0 / widthy_center = (ymin + ymax) / 2.0 / heightw = (xmax - xmin) / widthh = (ymax - ymin) / heightf.write(f"{cls_id} {x_center} {y_center} {w} {h}\n")# 遍歷VOC數據集的Annotations目錄,進行轉換
for voc_annotation in os.listdir(voc_annotations_path):if voc_annotation.endswith('.xml'):voc_annotation_file = os.path.join(voc_annotations_path, voc_annotation)image_id = os.path.splitext(voc_annotation)[0]voc_image_file = os.path.join(voc_images_path, f"{image_id}.jpg")yolo_label_file = os.path.join(yolo_labels_path, f"{image_id}.txt")yolo_image_file = os.path.join(yolo_images_path, f"{image_id}.jpg")convert_voc_to_yolo(voc_annotation_file, yolo_label_file)if os.path.exists(voc_image_file):shutil.copy(voc_image_file, yolo_image_file)print("轉換完成!")
二、YOLO格式數據集劃分(訓練、驗證、測試)
參考:https://docs.ultralytics.com/datasets/detect/#ultralytics-yolo-format
隨機將數據集按照0.7-0.2-0.1比例劃分為訓練、驗證、測試數據集。
注意,修改代碼中圖片的后綴,如果是.jpg,就把.png修改為.jpg。
最終結果,
2.1 版本1
用版本1劃分就行,也可以用版本2,版本3就不用了。版本1和版本2是兩種不同的組織方式都能訓練。版本1是官方的組織方法。
import os
import shutil
import randomdef make_yolo_dataset(images_folder, labels_folder, output_folder, train_ratio=0.8):# 創建目標文件夾images_train_folder = os.path.join(output_folder, 'images/train')images_val_folder = os.path.join(output_folder, 'images/val')labels_train_folder = os.path.join(output_folder, 'labels/train')labels_val_folder = os.path.join(output_folder, 'labels/val')os.makedirs(images_train_folder, exist_ok=True)os.makedirs(images_val_folder, exist_ok=True)os.makedirs(labels_train_folder, exist_ok=True)os.makedirs(labels_val_folder, exist_ok=True)# 獲取圖片和標簽的文件名(不包含擴展名)image_files = [f for f in os.listdir(images_folder) if f.endswith('.jpg')]label_files = [f for f in os.listdir(labels_folder) if f.endswith('.txt')]image_base_names = set(os.path.splitext(f)[0] for f in image_files)label_base_names = set(os.path.splitext(f)[0] for f in label_files)# 找出圖片和標簽都存在的文件名matched_files = list(image_base_names & label_base_names)# 打亂順序并劃分為訓練集和驗證集random.shuffle(matched_files)split_idx = int(len(matched_files) * train_ratio)train_files = matched_files[:split_idx]val_files = matched_files[split_idx:]# 移動文件到對應文件夾for base_name in train_files:img_src = os.path.join(images_folder, f"{base_name}.jpg")lbl_src = os.path.join(labels_folder, f"{base_name}.txt")img_dst = os.path.join(images_train_folder, f"{base_name}.jpg")lbl_dst = os.path.join(labels_train_folder, f"{base_name}.txt")shutil.copyfile(img_src, img_dst)shutil.copyfile(lbl_src, lbl_dst)for base_name in val_files:img_src = os.path.join(images_folder, f"{base_name}.jpg")lbl_src = os.path.join(labels_folder, f"{base_name}.txt")img_dst = os.path.join(images_val_folder, f"{base_name}.jpg")lbl_dst = os.path.join(labels_val_folder, f"{base_name}.txt")shutil.copyfile(img_src, img_dst)shutil.copyfile(lbl_src, lbl_dst)print("數據集劃分完成!")# 使用示例
images_folder = 'path/to/your/images_folder' # 原始圖片文件夾路徑
labels_folder = 'path/to/your/labels_folder' # 原始標簽文件夾路徑
output_folder = 'path/to/your/output_folder' # 存放結果數據集的文件夾路徑
make_yolo_dataset(images_folder, labels_folder, output_folder)
2.2 版本2
import os
import shutil
import random
from math import floor# 創建輸出目錄的函數
def create_dirs(output_dir):images_dir = os.path.join(output_dir, 'images')labels_dir = os.path.join(output_dir, 'labels')for split in ['train', 'val', 'test']:os.makedirs(os.path.join(images_dir, split), exist_ok=True)os.makedirs(os.path.join(labels_dir, split), exist_ok=True)return images_dir, labels_dir# 獲取圖片和對應txt標簽的列表
def get_files(images_path, labels_path):image_files = [f for f in os.listdir(images_path) if f.endswith(('jpg', 'png', 'jpeg'))]label_files = [f for f in os.listdir(labels_path) if f.endswith('.txt')]# 檢查圖片和標簽是否配對paired_files = []for image_file in image_files:base_name = os.path.splitext(image_file)[0]label_file = base_name + '.txt'if label_file in label_files:paired_files.append((image_file, label_file))return paired_files# 將文件按比例劃分并拷貝到相應目錄
def split_and_copy(paired_files, images_path, labels_path, images_dir, labels_dir, train_ratio, val_ratio):random.shuffle(paired_files) # 隨機打亂total_files = len(paired_files)train_count = floor(total_files * train_ratio)val_count = floor(total_files * val_ratio)test_count = total_files - train_count - val_countsplits = {'train': paired_files[:train_count],'val': paired_files[train_count:train_count + val_count],'test': paired_files[train_count + val_count:]}for split, files in splits.items():for image_file, label_file in files:shutil.copy(os.path.join(images_path, image_file), os.path.join(images_dir, split, image_file))shutil.copy(os.path.join(labels_path, label_file), os.path.join(labels_dir, split, label_file))print(f'{split}: {len(files)} files')# 主函數
def main():# 寫死的路徑images_path = "E:\\DataSet\\LC\\large_coal_blocked_yolo\\totalImages" # 替換為實際圖片文件夾路徑labels_path = "E:\\DataSet\\LC\\large_coal_blocked_yolo\\totalLabels" # 替換為實際txt文件夾路徑output_dir = "E:\\DataSet\\LC\\large_coal_blocked_yolo\\output" # 替換為實際輸出主目錄路徑# 數據劃分比例train_ratio = 0.7val_ratio = 0.3test_ratio = 0# 容差值用于浮點數比較epsilon = 1e-6# 確保比例之和等于1assert abs(train_ratio + val_ratio + test_ratio - 1) < epsilon, "比例之和必須等于1"# 創建目錄images_dir, labels_dir = create_dirs(output_dir)# 獲取文件列表paired_files = get_files(images_path, labels_path)# 進行劃分并拷貝split_and_copy(paired_files, images_path, labels_path, images_dir, labels_dir, train_ratio, val_ratio)# 調用主函數
if __name__ == "__main__":main()
2.3 版本3
import os
import shutil
import random# YOLO格式數據集保存路徑
yolo_images_path1 = 'E:\\DataSet\\helmet-VOC'
yolo_labels_path1 = 'E:\\DataSet\\helmet-YOLO'
yolo_data_path = yolo_labels_path1yolo_images_path = os.path.join(yolo_images_path1, 'JPEGImages')
yolo_labels_path = os.path.join(yolo_labels_path1, 'labels')# 創建劃分后的目錄結構
train_images_path = os.path.join(yolo_data_path, 'train', 'images')
train_labels_path = os.path.join(yolo_data_path, 'train', 'labels')
val_images_path = os.path.join(yolo_data_path, 'val', 'images')
val_labels_path = os.path.join(yolo_data_path, 'val', 'labels')
test_images_path = os.path.join(yolo_data_path, 'test', 'images')
test_labels_path = os.path.join(yolo_data_path, 'test', 'labels')os.makedirs(train_images_path, exist_ok=True)
os.makedirs(train_labels_path, exist_ok=True)
os.makedirs(val_images_path, exist_ok=True)
os.makedirs(val_labels_path, exist_ok=True)
os.makedirs(test_images_path, exist_ok=True)
os.makedirs(test_labels_path, exist_ok=True)# 獲取所有圖片文件名(不包含擴展名)
image_files = [f[:-4] for f in os.listdir(yolo_images_path) if f.endswith('.png')]# 隨機打亂文件順序
random.shuffle(image_files)# 劃分數據集比例
train_ratio = 0.7
val_ratio = 0.2
test_ratio = 0.1train_count = int(train_ratio * len(image_files))
val_count = int(val_ratio * len(image_files))
test_count = len(image_files) - train_count - val_counttrain_files = image_files[:train_count]
val_files = image_files[train_count:train_count + val_count]
test_files = image_files[train_count + val_count:]# 移動文件到相應的目錄
def move_files(files, src_images_path, src_labels_path, dst_images_path, dst_labels_path):for file in files:src_image_file = os.path.join(src_images_path, f"{file}.png")src_label_file = os.path.join(src_labels_path, f"{file}.txt")dst_image_file = os.path.join(dst_images_path, f"{file}.png")dst_label_file = os.path.join(dst_labels_path, f"{file}.txt")if os.path.exists(src_image_file) and os.path.exists(src_label_file):shutil.move(src_image_file, dst_image_file)shutil.move(src_label_file, dst_label_file)# 移動訓練集文件
move_files(train_files, yolo_images_path, yolo_labels_path, train_images_path, train_labels_path)
# 移動驗證集文件
move_files(val_files, yolo_images_path, yolo_labels_path, val_images_path, val_labels_path)
# 移動測試集文件
move_files(test_files, yolo_images_path, yolo_labels_path, test_images_path, test_labels_path)print("數據集劃分完成!")
三、一步到位
如果不想分兩步進行格式轉換,那么以下腳本結合了以上兩步,直接得到最后按比例劃分訓練、驗證、測試的數據集結果。
注意:需要修改 voc_data_path ,yolo_data_path ,class_mapping 以及 ‘.png’ 后綴。
import os
import shutil
import random
import xml.etree.ElementTree as ET
from tqdm import tqdm# VOC格式數據集路徑
voc_data_path = 'E:\\DataSet-VOC'
voc_annotations_path = os.path.join(voc_data_path, 'Annotations')
voc_images_path = os.path.join(voc_data_path, 'JPEGImages')# YOLO格式數據集保存路徑
yolo_data_path = 'E:\\DataSet-YOLO'
yolo_images_path = os.path.join(yolo_data_path, 'images')
yolo_labels_path = os.path.join(yolo_data_path, 'labels')# 創建YOLO格式數據集目錄
os.makedirs(yolo_images_path, exist_ok=True)
os.makedirs(yolo_labels_path, exist_ok=True)# 類別映射 (可以根據自己的數據集進行調整)
class_mapping = {'head': 0,'helmet': 1,'person': 2,# 添加更多類別...
}def convert_voc_to_yolo(voc_annotation_file, yolo_label_file):tree = ET.parse(voc_annotation_file)root = tree.getroot()size = root.find('size')width = float(size.find('width').text)height = float(size.find('height').text)with open(yolo_label_file, 'w') as f:for obj in root.findall('object'):cls = obj.find('name').textif cls not in class_mapping:continuecls_id = class_mapping[cls]xmlbox = obj.find('bndbox')xmin = float(xmlbox.find('xmin').text)ymin = float(xmlbox.find('ymin').text)xmax = float(xmlbox.find('xmax').text)ymax = float(xmlbox.find('ymax').text)x_center = (xmin + xmax) / 2.0 / widthy_center = (ymin + ymax) / 2.0 / heightw = (xmax - xmin) / widthh = (ymax - ymin) / heightf.write(f"{cls_id} {x_center} {y_center} {w} {h}\n")# 遍歷VOC數據集的Annotations目錄,進行轉換
print("開始VOC到YOLO格式轉換...")
for voc_annotation in tqdm(os.listdir(voc_annotations_path)):if voc_annotation.endswith('.xml'):voc_annotation_file = os.path.join(voc_annotations_path, voc_annotation)image_id = os.path.splitext(voc_annotation)[0]voc_image_file = os.path.join(voc_images_path, f"{image_id}.png")yolo_label_file = os.path.join(yolo_labels_path, f"{image_id}.txt")yolo_image_file = os.path.join(yolo_images_path, f"{image_id}.png")convert_voc_to_yolo(voc_annotation_file, yolo_label_file)if os.path.exists(voc_image_file):shutil.copy(voc_image_file, yolo_image_file)print("VOC到YOLO格式轉換完成!")# 劃分數據集
train_images_path = os.path.join(yolo_data_path, 'train', 'images')
train_labels_path = os.path.join(yolo_data_path, 'train', 'labels')
val_images_path = os.path.join(yolo_data_path, 'val', 'images')
val_labels_path = os.path.join(yolo_data_path, 'val', 'labels')
test_images_path = os.path.join(yolo_data_path, 'test', 'images')
test_labels_path = os.path.join(yolo_data_path, 'test', 'labels')os.makedirs(train_images_path, exist_ok=True)
os.makedirs(train_labels_path, exist_ok=True)
os.makedirs(val_images_path, exist_ok=True)
os.makedirs(val_labels_path, exist_ok=True)
os.makedirs(test_images_path, exist_ok=True)
os.makedirs(test_labels_path, exist_ok=True)# 獲取所有圖片文件名(不包含擴展名)
image_files = [f[:-4] for f in os.listdir(yolo_images_path) if f.endswith('.png')]# 隨機打亂文件順序
random.shuffle(image_files)# 劃分數據集比例
train_ratio = 0.7
val_ratio = 0.2
test_ratio = 0.1train_count = int(train_ratio * len(image_files))
val_count = int(val_ratio * len(image_files))
test_count = len(image_files) - train_count - val_counttrain_files = image_files[:train_count]
val_files = image_files[train_count:train_count + val_count]
test_files = image_files[train_count + val_count:]# 移動文件到相應的目錄
def move_files(files, src_images_path, src_labels_path, dst_images_path, dst_labels_path):for file in tqdm(files):src_image_file = os.path.join(src_images_path, f"{file}.png")src_label_file = os.path.join(src_labels_path, f"{file}.txt")dst_image_file = os.path.join(dst_images_path, f"{file}.png")dst_label_file = os.path.join(dst_labels_path, f"{file}.txt")if os.path.exists(src_image_file) and os.path.exists(src_label_file):shutil.move(src_image_file, dst_image_file)shutil.move(src_label_file, dst_label_file)# 移動訓練集文件
print("移動訓練集文件...")
move_files(train_files, yolo_images_path, yolo_labels_path, train_images_path, train_labels_path)
# 移動驗證集文件
print("移動驗證集文件...")
move_files(val_files, yolo_images_path, yolo_labels_path, val_images_path, val_labels_path)
# 移動測試集文件
print("移動測試集文件...")
move_files(test_files, yolo_images_path, yolo_labels_path, test_images_path, test_labels_path)print("數據集劃分完成!")# 刪除原始的 images 和 labels 文件夾
shutil.rmtree(yolo_images_path)
shutil.rmtree(yolo_labels_path)print("原始 images 和 labels 文件夾刪除完成!")