WIDER FACE数据集转YOLO格式
1. 引出问题
本人最近在做毕设相关内容，第一阶段目标是通过目标检测来统计课堂人数，因此需要对人脸和人头进行目标检测。模型方面没什么好说的无脑用YOLO，数据集方面，人脸部分找到了来自港中文的WIDER FACE数据集。但是解压后发现标注数据并不是YOLO格式的，因此通过分析后写了一个简单的脚本进行转换。
2. 标注数据转YOLO格式
2.1 源代码
通过分析发现原标注数据的格式如下：
[图片路径]
[该图片中检测框数量]
[左上角x坐标 左上角y坐标 检测框宽 检测框高 …(一些额外信息)]
…
因此得到如下格式转换代码：
import os
from PIL import Imageparent_path = "P:/Graduation_Design/Dataset/WIDER_FACE/"def convert_to_yolo_format(input_file, output_dir, image_dir):with open(input_file, 'r') as f:lines = f.readlines()i = 0while i < len(lines):image_path = lines[i].strip() # Get the relative path of imagenum_boxes = int(lines[i + 1].strip()) # Get the number of boxes# Path of the label filelabel_path = os.path.join(output_dir, os.path.basename(image_path).replace('.jpg', '.txt'))os.makedirs(os.path.dirname(label_path), exist_ok=True)# Get the Absolute Path of the imageimage_abs_path = os.path.join(image_dir, image_path)# Open the image to get the real size of itwith Image.open(image_abs_path) as img:img_width, img_height = img.size# Create the file and write data inwith open(label_path, 'w') as label_file:for j in range(num_boxes):# Fetch the box data (x_min, y_min, width, height)box_data = list(map(int, lines[i + 2 + j].strip().split()))x_min, y_min, width, height = box_data[:4]# Calculate the center coordinate (x_center, y_center)x_center = (x_min + width / 2)y_center = (y_min + height / 2)# Convert to the relative coordinatesx_center /= img_widthy_center /= img_heightwidth /= img_widthheight /= img_height# The class is defaulted by 0label_file.write(f"0 {x_center} {y_center} {width} {height}\n")# Update the index and jump to the next imagei += 2 + (1 if num_boxes == 0 else num_boxes)if __name__ == "__main__":# Modify the additional section by your own pathinput_path = parent_path+"wider_face_split/"output_path = parent_path+"wider_for_yolo/"input_file_pre = "wider_face_"input_file_sub = "_bbx_gt.txt"if not os.path.exists(output_path):os.makedirs(output_path)# Train and Validationdatasetfile = ["train", "val"]for category in datasetfile:convert_to_yolo_format(input_path + input_file_pre + category + input_file_sub,output_path + category + "/labels",parent_path+f"WIDER_{category}/WIDER_{category}/images")
2.2 使用方法
- 首先修改代码第五行的
parent_path变量，修改成自己解压后WIDER FACE数据集所在的路径。 - 接着在主程序部分将
input_path改成自己下载的标注数据所在文件夹路径 - 随后在主程序部分将
output_path改成自己需要导出YOLO格式标注数据的路径
3. 合并图片到一个目录
标注数据格式转换完成后，由于原数据集文件夹中对不同情景下的图片做了一个分类，而YOLO训练时要求所有训练图片在一个文件夹下，因此训练前还需要将数据集所有图片copy到一个文件夹下。
3.1 源代码
import os
import shutildef copy_images(src_dir, dest_dir):# 确保目标目录存在if not os.path.exists(dest_dir):os.makedirs(dest_dir)# 递归查找所有图片for root, _, files in os.walk(src_dir):for file in files:if file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp')):src_path = os.path.join(root, file)dest_path = os.path.join(dest_dir, file)# 如果目标文件已存在，可以选择覆盖或跳过if not os.path.exists(dest_path):shutil.copy2(src_path, dest_path) # 保留原文件元数据print(f"Copied: {src_path} -> {dest_path}")else:print(f"Skipped (already exists): {dest_path}")# 配置源文件夹和目标文件夹路径
train_source_folder = r"P:\Graduation_Design\Dataset\WIDER_FACE\WIDER_train\WIDER_train\images"
train_destination_folder = r"P:\Graduation_Design\Dataset\WIDER_FACE\WIDER_train\WIDER_train\data"
val_source_folder = r"P:\Graduation_Design\Dataset\WIDER_FACE\WIDER_val\WIDER_val\images"
val_destination_folder = r"P:\Graduation_Design\Dataset\WIDER_FACE\WIDER_val\WIDER_val\data"# 执行复制
copy_images(train_source_folder, train_destination_folder)
copy_images(val_source_folder, val_destination_folder)
3.2 使用方法
- 首先将
train_source_folder和val_source_folder修改为自己解压后的WIDER FACE数据集的训练集和验证集的images路径 - 接着将
train_destination_folder和val_destination_folder修改为自己导出转换后的标注数据的路径
4. 使用YOLO模型进行训练
经过上面的标注数据格式转换和图片合并两个步骤后，确保你得到了train和val两个路径：
并且这两个文件夹中的结构应该都如下图所示：
这里的classes.txt只需要写一个Face或者人脸即可
一切准备就绪后建立一个data.yaml文件填入如下内容：
train: P:/Graduation_Design/Dataset/WIDER_FACE/wider_for_yolo/train # 这里改成自己处理后的训练集路径
val: P:/Graduation_Design/Dataset/WIDER_FACE/wider_for_yolo/val # 这里改成自己处理后的验证集路径nc: 1 # 类别数
names: ['face'] # 类别名称
至此便可以通过加载这个data.yaml文件完成YOLO模型的训练了。