SODA-数据集处理成适合YOLOV8训练的数据集

import os
import json
from PIL import Image


def open_annotations_json(image_name):
    # 构造标注文件名
    annotation_name = image_name.replace(".jpg", ".json")
    # 构造标注文件路径

    # 在 train、val 和 test 文件夹中查找标注文件
    for folder in ['train', 'val', 'test']:
        annotation_file_path = os.path.join("Annotations", folder, annotation_name)
        if os.path.exists(annotation_file_path):
            # 如果找到标注文件，则打开并处理
            with open(annotation_file_path, "r") as file:
                data = json.load(file)
                print(f"找到{annotation_file_path}对应的标注文件")
                # 在这里对标注文件进行处理
                return data  # 返回标注文件数据
            break
    else:
        print(f"未找到{annotation_file_path}对应的标注文件")


def is_points_inside_rect(points, left, upper, right, lower):
    for i in range(0, len(points), 2):
        x = points[i]
        y = points[i + 1]
        if x < left or x > right or y < upper or y > lower:
            return False
    return True


def get_polys_in_sub_images(left, upper, right, lower, annotations, crop_size):
    outcome = []
    for annotation in annotations:
        points = annotation['poly']
        category_id = annotation['category_id']

        # 判断多边形是否在子图像内
        flag = is_points_inside_rect(points, left, upper, right, lower)

        if flag:
            # 调整坐标，使左上角的点成为 (0, 0)
            adjusted_points = [category_id]
            for i in range(0, len(points), 2):
                x = points[i] - left
                y = points[i + 1] - upper
                adjusted_points.extend([x / crop_size, y / crop_size])

            outcome.append(adjusted_points)

    return outcome


def process_one_images(image_name, save_image_path, save_txt_path):
    data = open_annotations_json(image_name)
    if data:
        # 加载图片基本信息
        file_name = data['images']['file_name']
        height = data['images']['height']
        width = data['images']['width']
        print(f"图片文件名: {file_name}")
        print(f"图片高度: {height}")
        print(f"图片宽度: {width}")
    else:
        print("未找到对应的标注文件")
        return

    crop_size = 640

    image_name = os.path.join('Images', image_name)

    # 打开图像
    image = Image.open(image_name)

    # 计算裁剪的行数和列数
    rows = height // crop_size
    cols = width // crop_size

    # 创建保存路径
    if not os.path.exists(save_image_path):
        os.makedirs(save_image_path)
    if not os.path.exists(save_txt_path):
        os.makedirs(save_txt_path)

    # 开始裁剪和保存
    for row in range(rows):
        for col in range(cols):
            # 计算裁剪区域的坐标
            left = col * crop_size
            upper = row * crop_size
            right = left + crop_size
            lower = upper + crop_size

            tem_polys = get_polys_in_sub_images(left, upper, right, lower, data['annotations'], crop_size)
            if len(tem_polys) > 5:
                # 裁剪图像
                cropped_image = image.crop((left, upper, right, lower))
                # 构造保存路径
                save_filename = f"{file_name[:-4]}_crop_{row}_{col}.jpg"  # 生成保存文件名，例如：00001_crop_0_0.jpg
                save_file_path = os.path.join(save_image_path, save_filename)

                # 保存裁剪后的图像
                cropped_image.save(save_file_path)

                print(f"已保存裁剪后的图像: {save_file_path}")

                # 构造保存路径
                save_filename = f"{file_name[:-4]}_crop_{row}_{col}.txt"  # 生成保存文件名，例如：00001_crop_0_0.txt
                save_file_path = os.path.join(save_txt_path, save_filename)

                # 打开文件并写入数据
                with open(save_file_path, "w") as f:
                    for poly in tem_polys:
                        line = " ".join(map(str, poly)) + "\n"  # 将列表中的元素转换为字符串并用空格连接
                        f.write(line)

                print(f"已保存裁剪后的文本文件: {save_file_path}")


if __name__ == '__main__':
    # 获取图像文件夹中的所有文件
    image_files = os.listdir('Images')

    # 定义数据集划分比例
    val_ratio = 0.2
    train_ratio = 0.7
    test_ratio = 0.1

    # 计算各数据集的数量
    num_val = int(len(image_files) * val_ratio)
    num_train = int(len(image_files) * train_ratio)
    num_test = int(len(image_files) * test_ratio)

    # 打印数据集数量
    print(f"总共有 {len(image_files)} 张图片。")
    print(f"验证集数量：{num_val}，训练集数量：{num_train}，测试集数量：{num_test}。")

    # 首先创建保存图像和标签的文件夹
    save_dirs = ["val_images", "train_images", "test_images", "val_labels", "train_labels", "test_labels"]
    for save_dir in save_dirs:
        os.makedirs(save_dir, exist_ok=True)

    # 分配图像到不同的数据集文件夹
    val_images = image_files[:num_val]
    train_images = image_files[num_val:num_val + num_train]
    test_images = image_files[num_val + num_train:]

    # 遍历并处理验证集图像
    for image_name in val_images:
        save_image_path = "val_images"  # 保存验证集图像的路径
        save_txt_path = "val_labels"  # 保存验证集标签的路径
        process_one_images(image_name, save_image_path, save_txt_path)

    # 遍历并处理训练集图像
    for image_name in train_images:
        save_image_path = "train_images"  # 保存训练集图像的路径
        save_txt_path = "train_labels"  # 保存训练集标签的路径
        process_one_images(image_name, save_image_path, save_txt_path)

    # 遍历并处理测试集图像
    for image_name in test_images:
        save_image_path = "test_images"  # 保存测试集图像的路径
        save_txt_path = "test_labels"  # 保存测试集标签的路径
        process_one_images(image_name, save_image_path, save_txt_path)