hippocampus-segmentation/split_prepare.py at main · Lisa-suger/hippocampus-segmentation · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87

from PIL import Image
import os
import shutil
import random


def split_data(images_source_folder, labels_source_folder, train_folder, test_folder, test_ratio=0.18,
                train_ratio=0.8):
    # 确保输出文件夹存在
    os.makedirs(train_folder, exist_ok=True)
    os.makedirs(test_folder, exist_ok=True)

    # 为训练和测试集创建子文件夹
    os.makedirs(os.path.join(train_folder, 'train_img'), exist_ok=True)
    os.makedirs(os.path.join(train_folder, 'train_label'), exist_ok=True)
    os.makedirs(os.path.join(test_folder, 'test_img'), exist_ok=True)
    os.makedirs(os.path.join(test_folder, 'test_label'), exist_ok=True)

    # 获取源文件夹中所有文件
    images_name = os.listdir(images_source_folder)
    labels_name = os.listdir(labels_source_folder)
    images_name.sort()
    labels_name.sort()
    num = list(range(len(os.listdir(images_source_folder))))
    random.shuffle(num)
    # 打乱文件顺序以确保随机划分

    test_size = int(len(num) * test_ratio)
    train_files, train_label_files = zip(*[(images_name[i], labels_name[i]) for i in num[test_size:]])
    test_files, test_label_files = zip(*[(images_name[i], labels_name[i]) for i in num[:test_size]])

    # 划分和复制文件到相应文件夹
    for file in train_files+ test_files:
        if file.lower().endswith('.nii.gz'):
            img_source_path = os.path.join(images_source_folder, file)

            if 'preprocessed_' in file.lower():
                if file in train_files:
                    dest_folder = 'train_img'
                elif file in test_files:
                    dest_folder = 'test_img'
            else:
                continue  # 不是标签，跳过处理

            # 构造目标路径
            base_name = os.path.splitext(file)[0]
            if file in train_files:
                dest_path = os.path.join(train_folder, dest_folder,file)
            elif file in test_files:
                dest_path = os.path.join(test_folder, dest_folder, file)

            # 复制文件到目标路径（这里假设文件已经被调整大小）
            shutil.copy(img_source_path, dest_path)

    for file in train_label_files  + test_label_files:
        if file.lower().endswith('.nii.gz'):
            label_source_path = os.path.join(labels_source_folder, file)

            if "mask" in file.lower():
                if file in train_label_files:
                    dest_folder = 'train_label'
                elif file in test_label_files:
                    dest_folder = 'test_label'
            else:
                continue  # 不是标签，跳过处理

            # 构造目标路径
            base_name = os.path.splitext(file)[0]
            if file in train_label_files:
                dest_path = os.path.join(train_folder, dest_folder, file)
            elif file in test_label_files:
                dest_path = os.path.join(test_folder, dest_folder, file)

            # 复制文件到目标路径（这里假设文件已经被调整大小）
            shutil.copy(label_source_path, dest_path)
    print("已经划分好数据集\n"
          "训练集||测试集\n"
          "0.82 || 0.18")


if __name__ == '__main__':
    split_data("initial_data/prepared_datas/prepared_images",
               "initial_data/prepared_datas/prepared_labels",
               "initial_data/data/train_data",
               "initial_data/data/test_data"
               )