WIDER FACE数据集转YOLO格式

1. 引出问题

本人最近在做毕设相关内容,第一阶段目标是通过目标检测来统计课堂人数,因此需要对人脸和人头进行目标检测。模型方面没什么好说的无脑用YOLO,数据集方面,人脸部分找到了来自港中文的WIDER FACE数据集。但是解压后发现标注数据并不是YOLO格式的,因此通过分析后写了一个简单的脚本进行转换。

2. 标注数据转YOLO格式

2.1 源代码

通过分析发现原标注数据的格式如下:

[图片路径]

[该图片中检测框数量]

[左上角x坐标 左上角y坐标 检测框宽 检测框高 ...(一些额外信息)]

...

因此得到如下格式转换代码:

python 复制代码
import os
from PIL import Image


parent_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 = 0
    while i < len(lines):
        image_path = lines[i].strip()  # Get the relative path of image
        num_boxes = int(lines[i + 1].strip())  # Get the number of boxes
        # Path of the label file
        label_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 image
        image_abs_path = os.path.join(image_dir, image_path)

        # Open the image to get the real size of it
        with Image.open(image_abs_path) as img:
            img_width, img_height = img.size

        # Create the file and write data in
        with 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 coordinates
                x_center /= img_width
                y_center /= img_height
                width /= img_width
                height /= img_height

                # The class is defaulted by 0
                label_file.write(f"0 {x_center} {y_center} {width} {height}\n")

        # Update the index and jump to the next image
        i += 2 + (1 if num_boxes == 0 else num_boxes)


if __name__ == "__main__":
    # Modify the additional section by your own path
    input_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 Validation
    datasetfile = ["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 源代码

python 复制代码
import os
import shutil


def 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_folderval_source_folder修改为自己解压后的WIDER FACE数据集的训练集和验证集的images路径
  • 接着将train_destination_folderval_destination_folder修改为自己导出转换后的标注数据的路径

4. 使用YOLO模型进行训练

经过上面的标注数据格式转换和图片合并两个步骤后,确保你得到了trainval两个路径:

并且这两个文件夹中的结构应该都如下图所示:

这里的classes.txt只需要写一个Face或者人脸即可

一切准备就绪后建立一个data.yaml文件填入如下内容:

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模型的训练了。

相关推荐
董董灿是个攻城狮1 小时前
020:为什么 Resnet 如此重要?
人工智能·计算机视觉·cnn
汪子熙2 小时前
为什么 BERT 仅使用 Transformer 的编码器部分,而不使用解码器部分?
人工智能
AI服务老曹2 小时前
满足不同场景的需求的智慧物流开源了
人工智能·开源·自动化·能源
程序猿阿伟2 小时前
《开源与合作:驱动鸿蒙Next系统中人工智能技术创新发展的双引擎》
人工智能·开源·harmonyos
pchmi2 小时前
C# OpenCV机器视觉:常用滤波算法
人工智能·opencv·c#·机器视觉·中值滤波
茫茫人海一粒沙2 小时前
稀疏检索、密集检索与混合检索:概念、技术与应用
人工智能
代数狂人2 小时前
NPC与AI深度融合结合雷鸟X3Pro AR智能眼镜:引领游戏行业沉浸式与增强现实新纪元的畅想
人工智能·游戏·ar
卧式纯绿3 小时前
自动驾驶3D目标检测综述(八)
人工智能·目标检测·3d·目标跟踪·cnn·自动驾驶
AIGC大时代3 小时前
数据分析如何正确使用ChatGPT进行辅助?
大数据·人工智能·深度学习·chatgpt·数据挖掘·数据分析·aigc
懒麻蛇3 小时前
ChatGPT Task功能初探
人工智能·chatgpt