本文nnunetv2实验中BraTS数据集转换成MSD数据集格式的代码, nnunetv1的dataset.json与nnunetv2略有不同

BraTS数据集转换成MSD数据集#

BraTS2019:#

1
import os
2
import shutil
3
import json
4
import nibabel as nib
5
import numpy as np
6
from collections import OrderedDict
7

8
def convert_brats2019_to_nnunet(brats_root, nnunet_raw_data_base):
9
    """
10
    将BraTS2019数据集转换为nnUNet格式，处理标签映射
11

12
    Args:
13
        brats_root: BraTS2019原始数据根目录路径
14
        nnunet_raw_data_base: nnUNet原始数据基础目录路径
15
    """
16

17
    # 设置路径
18
    task_name = "Dataset001_BraTS2019"
19
    task_folder = os.path.join(nnunet_raw_data_base, "nnUNet_raw_data", task_name)
20

21
    # 创建必要的目录
22
    imagesTr_folder = os.path.join(task_folder, "imagesTr")
23
    imagesTs_folder = os.path.join(task_folder, "imagesTs")
24
    labelsTr_folder = os.path.join(task_folder, "labelsTr")
25
    labelsTs_folder = os.path.join(task_folder, "labelsTs")
26

27
    for folder in [imagesTr_folder, imagesTs_folder, labelsTr_folder, labelsTs_folder]:
28
        os.makedirs(folder, exist_ok=True)
29

30
    # 获取训练数据和测试数据路径
31
    hgg_folder = os.path.join(brats_root, "HGG")
32
    lgg_folder = os.path.join(brats_root, "LGG")
33

34
    # 收集所有训练病例
35
    train_cases = []
36

37
    # 处理HGG数据
38
    if os.path.exists(hgg_folder):
39
        hgg_cases = [d for d in os.listdir(hgg_folder) if os.path.isdir(os.path.join(hgg_folder, d))]
40
        for case in hgg_cases:
41
            train_cases.append(("HGG", case))
42

43
    # 处理LGG数据
44
    if os.path.exists(lgg_folder):
45
        lgg_cases = [d for d in os.listdir(lgg_folder) if os.path.isdir(os.path.join(lgg_folder, d))]
46
        for case in lgg_cases:
47
            train_cases.append(("LGG", case))
48

49
    print(f"找到 {len(train_cases)} 个训练病例")
50

51
    # 模态映射
52
    modality_mapping = {
53
        'flair': '0000',
54
        't1ce': '0001',
55
        't1': '0002',
56
        't2': '0003'
57
    }
58

59
    def process_label(label_path, output_path):
60
        """
61
        处理BraTS标签，将原始标签映射为连续的标签
62
        BraTS原始标签: 0(背景), 1(坏死/非增强肿瘤), 2(水肿), 4(增强肿瘤)
63
        映射后标签: 0(背景), 1(坏死/非增强肿瘤), 2(水肿), 3(增强肿瘤)
64
        """
65
        # 加载标签数据
66
        label_nii = nib.load(label_path)
67
        label_data = label_nii.get_fdata().astype(np.uint8)
68

69
        # 检查原始标签值
70
        unique_labels = np.unique(label_data)
71
        print(f"处理 {os.path.basename(label_path)}，原始标签值: {unique_labels}")
72

73
        # 创建新的标签数组
74
        new_label_data = np.zeros_like(label_data)
75

76
        # 标签映射: 0->0, 1->1, 2->2, 4->3
77
        new_label_data[label_data == 0] = 0  # 背景
78
        new_label_data[label_data == 1] = 1  # 坏死和非增强肿瘤
79
        new_label_data[label_data == 2] = 2  # 水肿
80
        new_label_data[label_data == 4] = 3  # 增强肿瘤
81

82
        # 检查映射后的标签值
83
        new_unique_labels = np.unique(new_label_data)
84
        print(f"映射后标签值: {new_unique_labels}")
85

86
        # 保存新的标签文件
87
        new_label_nii = nib.Nifti1Image(new_label_data, label_nii.affine, label_nii.header)
88
        nib.save(new_label_nii, output_path)
89

90
    training_cases = []
91
    test_cases = []
92

93
    for i, (grade, case_name) in enumerate(train_cases):
94
        case_folder = os.path.join(brats_root, grade, case_name)
95

96
        if not os.path.exists(case_folder):
97
            print(f"警告: 病例文件夹不存在 {case_folder}")
98
            continue
99

100
        # 检查所有必需的文件是否存在
101
        required_files = {
102
            'flair': f"{case_name}_flair.nii",
103
            't1ce': f"{case_name}_t1ce.nii",
104
            't1': f"{case_name}_t1.nii",
105
            't2': f"{case_name}_t2.nii",
106
            'seg': f"{case_name}_seg.nii"
107
        }
108

109
        # 检查.nii.gz格式
110
        for key, filename in required_files.items():
111
            if not os.path.exists(os.path.join(case_folder, filename)):
112
                # 尝试.nii.gz格式
113
                gz_filename = filename + ".gz"
114
                if os.path.exists(os.path.join(case_folder, gz_filename)):
115
                    required_files[key] = gz_filename
116
                else:
117
                    print(f"警告: 文件 {filename} 或 {gz_filename} 不存在于 {case_folder}")
118
                    break
119
        else:
120
            # 所有文件都存在，处理这个病例
121

122
            # 决定这个病例是用于训练还是测试（这里简单地将前80%用于训练）
123
            if i < len(train_cases) * 0.8:
124
                # 训练数据
125
                for modality, suffix in modality_mapping.items():
126
                    src_file = os.path.join(case_folder, required_files[modality])
127
                    dst_file = os.path.join(imagesTr_folder, f"{case_name}_{suffix}.nii.gz")
128

129
                    # 如果源文件不是.gz格式，需要压缩
130
                    if not src_file.endswith('.gz'):
131
                        img = nib.load(src_file)
132
                        nib.save(img, dst_file)
133
                    else:
134
                        shutil.copy2(src_file, dst_file)
135

136
                # 处理分割标签（重要：进行标签映射）
137
                src_seg = os.path.join(case_folder, required_files['seg'])
138
                dst_seg = os.path.join(labelsTr_folder, f"{case_name}.nii.gz")
139
                process_label(src_seg, dst_seg)
140

141
                training_cases.append(case_name)
142

143
            else:
144
                # 测试数据
145
                for modality, suffix in modality_mapping.items():
146
                    src_file = os.path.join(case_folder, required_files[modality])
147
                    dst_file = os.path.join(imagesTs_folder, f"{case_name}_{suffix}.nii.gz")
148

149
                    if not src_file.endswith('.gz'):
150
                        img = nib.load(src_file)
151
                        nib.save(img, dst_file)
152
                    else:
153
                        shutil.copy2(src_file, dst_file)
154

155
                # 测试数据的标签（同样进行标签映射）
156
                src_seg = os.path.join(case_folder, required_files['seg'])
157
                dst_seg = os.path.join(labelsTs_folder, f"{case_name}.nii.gz")
158
                process_label(src_seg, dst_seg)
159

160
                test_cases.append(case_name)
161

162
    print(f"处理完成: {len(training_cases)} 个训练病例, {len(test_cases)} 个测试病例")
163

164
    # 创建dataset.json文件
165
    dataset_json = OrderedDict()
166
    dataset_json['name'] = "BraTS2019"
167
    dataset_json['description'] = "Brain Tumor Segmentation Challenge 2019"
168
    dataset_json['tensorImageSize'] = "4D"
169
    dataset_json['reference'] = "https://www.med.upenn.edu/cbica/brats2019/"
170
    dataset_json['licence'] = "see BraTS2019 website"
171
    dataset_json['release'] = "1.0"
172

173
    # 模态信息
174
    dataset_json['channel_names'] = {
175
        "0": "FLAIR",
176
        "1": "T1ce",
177
        "2": "T1",
178
        "3": "T2"
179
    }
180

181
    # 标签信息 - 修正为包含所有必要的标签
182
    dataset_json['labels'] = {
183
        "background": 0,
184
        "necrotic/non-enhancing tumor": 1,
185
        "edema": 2,
186
        "enhancing tumor": 3
187
    }
188
    dataset_json['file_ending'] = ".nii.gz"   #按需修改 .nii或.nii.gz
189
    # 训练和测试数据列表
190
    dataset_json['numTraining'] = len(training_cases)
191
    dataset_json['numTest'] = len(test_cases)
192

193
    dataset_json['training'] = []
194
    for case in training_cases:
195
        case_dict = {
196
            "image": f"./imagesTr/{case}.nii.gz",
197
            "label": f"./labelsTr/{case}.nii.gz"
198
        }
199
        dataset_json['training'].append(case_dict)
200

201
    dataset_json['test'] = []
202
    for case in test_cases:
203
        dataset_json['test'].append(f"./imagesTs/{case}.nii.gz")
204

205
    # 保存dataset.json
206
    json_file_path = os.path.join(task_folder, "dataset.json")
207
    with open(json_file_path, 'w') as f:
208
        json.dump(dataset_json, f, indent=4)
209

210
    print(f"dataset.json 已保存到: {json_file_path}")
211
    print("数据转换完成！标签已正确映射：0(背景) -> 0, 1(坏死) -> 1, 2(水肿) -> 2, 4(增强肿瘤) -> 3")
212

213
    return task_folder
214

215
# 使用示例
216
if __name__ == "__main__":
217
    # 设置路径
218
    brats_root = "/root/autodl-tmp/nnUNet_raw_data_base/BraTS2019"  # 您的BraTS2019数据根目录
219
    nnunet_raw_data_base = "/root/autodl-tmp/nnUNet_raw_data_base"  # nnUNet原始数据基础目录
220

221
    # 执行转换
222
    convert_brats2019_to_nnunet(brats_root, nnunet_raw_data_base)

BraTS2023#

1
import os
2
import json
3
import shutil
4
import argparse
5
import threading
6
from collections import OrderedDict
7
from concurrent.futures import ThreadPoolExecutor, as_completed
8
from datetime import datetime
9

10
import numpy as np
11
import nibabel as nib
12

13
try:
14
    from tqdm import tqdm
15
except ImportError:  # 没装 tqdm 也能跑
16
    tqdm = None
17

18

19
# BraTS2023 模态命名 -> nnU-Net 通道号
20
MODALITY_MAPPING = OrderedDict([
21
    ('t1n', '0000'),   # T1 native
22
    ('t1c', '0001'),   # T1 contrast enhanced
23
    ('t2f', '0002'),   # T2 FLAIR
24
    ('t2w', '0003'),   # T2 weighted
25
])
26

27
_print_lock = threading.Lock()
28

29

30
def log(msg):
31
    """线程安全打印。"""
32
    with _print_lock:
33
        print(msg, flush=True)
34

35

36
def check_file_validity(file_path):
37
    """检查文件是否存在、非空且 NIfTI 头可读（只读头，很快）。"""
38
    try:
39
        if not os.path.exists(file_path):
40
            return False, "文件不存在"
41
        if os.path.getsize(file_path) == 0:
42
            return False, "文件为空"
43
        nib.load(file_path)          # 只解析 header，不读数据
44
        return True, "文件有效"
45
    except Exception as e:
46
        return False, f"文件无效: {str(e)}"
47

48

49
def resolve_required_files(case_folder, base_name):
50
    """定位一个病例的 5 个文件（4 模态 + seg），同时支持 .nii / .nii.gz。"""
51
    files, errors = {}, []
52
    for key in list(MODALITY_MAPPING.keys()) + ['seg']:
53
        fname = f"{base_name}-{key}.nii"
54
        plain = os.path.join(case_folder, fname)
55
        gz = plain + ".gz"
56

57
        if os.path.exists(gz):
58
            chosen = gz
59
        elif os.path.exists(plain):
60
            chosen = plain
61
        else:
62
            errors.append(f"文件缺失 - 病例: {base_name}, 文件: {fname} 或 {fname}.gz")
63
            continue
64

65
        ok, msg = check_file_validity(chosen)
66
        if not ok:
67
            errors.append(f"文件无效 - 病例: {base_name}, {os.path.basename(chosen)}: {msg}")
68
            continue
69
        files[key] = chosen
70
    return files, errors
71

72

73
def copy_to_gz(src, dst, overwrite=False):
74
    """把源文件落到目标 .nii.gz；非 .gz 源用 nibabel 重存为 .gz。写临时文件再原子改名。"""
75
    if not overwrite and os.path.exists(dst) and os.path.getsize(dst) > 0:
76
        return False  # 已存在，跳过（断点续跑）
77

78
    tmp = dst + ".tmp.nii.gz"
79
    try:
80
        if src.endswith(".gz"):
81
            shutil.copy2(src, tmp)
82
        else:
83
            nib.save(nib.load(src), tmp)
84
        os.replace(tmp, dst)
85
    finally:
86
        if os.path.exists(tmp):
87
            try:
88
                os.remove(tmp)
89
            except OSError:
90
                pass
91
    return True
92

93

94
def process_case(case_name, brats_root, img_dir, lbl_dir, split,
95
                 verify_labels=False, overwrite=False):
96
    """单个病例的完整处理，纯函数式：所有错误/日志放进返回值，不碰全局状态。"""
97
    result = {
98
        "case": case_name,
99
        "split": split,
100
        "ok": False,
101
        "errors": [],
102
        "label_values": None,
103
    }
104

105
    case_folder = os.path.join(brats_root, case_name)
106
    if not os.path.isdir(case_folder):
107
        result["errors"].append(f"病例文件夹不存在: {case_folder}")
108
        return result
109

110
    files, errors = resolve_required_files(case_folder, case_name)
111
    if errors:
112
        result["errors"].extend(errors)
113
        return result
114

115
    try:
116
        for modality, suffix in MODALITY_MAPPING.items():
117
            copy_to_gz(files[modality],
118
                       os.path.join(img_dir, f"{case_name}_{suffix}.nii.gz"),
119
                       overwrite)
120

121
        dst_seg = os.path.join(lbl_dir, f"{case_name}.nii.gz")
122
        copy_to_gz(files['seg'], dst_seg, overwrite)
123

124
        if verify_labels:
125
            data = np.asanyarray(nib.load(dst_seg).dataobj)
126
            result["label_values"] = sorted(int(v) for v in np.unique(data))
127

128
        result["ok"] = True
129
    except Exception as e:
130
        result["errors"].append(f"处理失败 - 病例: {case_name}, 错误: {str(e)}")
131

132
    return result
133

134

135
def convert_brats2023_to_nnunet(brats_root, nnunet_raw_data_base,
136
                                task_name="Dataset001_BraTS2023",
137
                                workers=8, train_ratio=0.8,
138
                                verify_labels=False, overwrite=False):
139
    task_folder = os.path.join(nnunet_raw_data_base, task_name)
140
    imagesTr = os.path.join(task_folder, "imagesTr")
141
    imagesTs = os.path.join(task_folder, "imagesTs")
142
    labelsTr = os.path.join(task_folder, "labelsTr")
143
    labelsTs = os.path.join(task_folder, "labelsTs")
144
    for folder in (imagesTr, imagesTs, labelsTr, labelsTs):
145
        os.makedirs(folder, exist_ok=True)
146

147
    # 收集病例并排序，保证划分可复现
148
    all_cases = sorted(
149
        item for item in os.listdir(brats_root)
150
        if item.startswith('BraTS-GLI-') and os.path.isdir(os.path.join(brats_root, item))
151
    )
152
    print(f"找到 {len(all_cases)} 个病例，使用 {workers} 个线程处理")
153
    if not all_cases:
154
        print("警告: 没有找到任何病例")
155
        return None
156

157
    # ---- 先确定划分，再并行 ----
158
    n_train = int(len(all_cases) * train_ratio)
159
    jobs = []
160
    for i, case_name in enumerate(all_cases):
161
        if i < n_train:
162
            jobs.append((case_name, "train", imagesTr, labelsTr))
163
        else:
164
            jobs.append((case_name, "test", imagesTs, labelsTs))
165

166
    results = []
167
    t0 = datetime.now()
168
    with ThreadPoolExecutor(max_workers=workers) as executor:
169
        futures = {
170
            executor.submit(process_case, case_name, brats_root, img_dir, lbl_dir,
171
                            split, verify_labels, overwrite): case_name
172
            for case_name, split, img_dir, lbl_dir in jobs
173
        }
174

175
        iterator = as_completed(futures)
176
        if tqdm is not None:
177
            iterator = tqdm(iterator, total=len(futures), desc="转换中", ncols=80)
178

179
        done = 0
180
        for future in iterator:
181
            case_name = futures[future]
182
            try:
183
                res = future.result()
184
            except Exception as e:  # 兜底，理论上 process_case 已经吞掉异常
185
                res = {"case": case_name, "split": "unknown", "ok": False,
186
                       "errors": [f"线程异常 - 病例: {case_name}, 错误: {str(e)}"],
187
                       "label_values": None}
188
            results.append(res)
189

190
            done += 1
191
            if tqdm is None and done % 50 == 0:
192
                log(f"进度: {done}/{len(futures)}")
193
            if not res["ok"]:
194
                log(f"[跳过] {case_name}: {res['errors'][0] if res['errors'] else '未知原因'}")
195

196
    elapsed = (datetime.now() - t0).total_seconds()
197

198
    # ---- 主线程统一汇总 ----
199
    results.sort(key=lambda r: r["case"])
200
    training_cases = [r["case"] for r in results if r["ok"] and r["split"] == "train"]
201
    test_cases = [r["case"] for r in results if r["ok"] and r["split"] == "test"]
202
    skipped_cases = [r["case"] for r in results if not r["ok"]]
203
    errors = [e for r in results for e in r["errors"]]
204

205
    if verify_labels:
206
        seen = sorted({v for r in results if r["label_values"] for v in r["label_values"]})
207
        print(f"数据集中出现的标签值: {seen}")
208

209
    print(f"\n处理完成: {len(training_cases)} 个训练病例, {len(test_cases)} 个测试病例")
210
    print(f"跳过的病例数量: {len(skipped_cases)}，耗时 {elapsed:.1f} 秒")
211

212
    # ---- 错误日志 ----
213
    error_file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "error.txt")
214
    with open(error_file_path, 'w', encoding='utf-8') as f:
215
        f.write("BraTS2023数据转换错误报告\n")
216
        f.write(f"生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
217
        f.write("=" * 80 + "\n\n")
218
        f.write(f"总计处理: {len(all_cases)} 个病例\n")
219
        f.write(f"成功处理: {len(training_cases) + len(test_cases)} 个病例\n")
220
        f.write(f"跳过病例: {len(skipped_cases)} 个病例\n")
221
        f.write(f"错误数量: {len(errors)} 个错误\n")
222
        f.write(f"总耗时: {elapsed:.1f} 秒 (线程数 {workers})\n\n")
223

224
        if skipped_cases:
225
            f.write("跳过的病例列表:\n")
226
            for case in skipped_cases:
227
                f.write(f"  - {case}\n")
228
            f.write("\n")
229

230
        if errors:
231
            f.write("详细错误信息:\n")
232
            for i, error in enumerate(errors, 1):
233
                f.write(f"{i}. {error}\n")
234
        else:
235
            f.write("没有发现错误。\n")
236
    print(f"错误日志已保存到: {error_file_path}")
237

238
    if not (training_cases or test_cases):
239
        print("警告: 没有成功处理任何病例，未生成dataset.json文件")
240
        return None
241

242
    # ---- dataset.json ----
243
    dataset_json = OrderedDict()
244
    dataset_json['name'] = "BraTS2023"
245
    dataset_json['description'] = "Brain Tumor Segmentation Challenge 2023"
246
    dataset_json['tensorImageSize'] = "4D"
247
    dataset_json['reference'] = "https://www.synapse.org/#!Synapse:syn51156910"
248
    dataset_json['licence'] = "see BraTS2023 website"
249
    dataset_json['release'] = "1.0"
250
    dataset_json['channel_names'] = {
251
        "0": "T1n",
252
        "1": "T1c",
253
        "2": "T2f",
254
        "3": "T2w",
255
    }
256
    dataset_json['labels'] = {
257
        "background": 0,
258
        "necrotic/non-enhancing tumor": 1,
259
        "edema": 2,
260
        "enhancing tumor": 3,
261
    }
262
    dataset_json["file_ending"] = ".nii.gz"
263
    dataset_json['numTraining'] = len(training_cases)
264
    dataset_json['numTest'] = len(test_cases)
265

266
    dataset_json['training'] = [
267
        {
268
            "image": [f"./imagesTr/{case}_{suffix}.nii.gz" for suffix in MODALITY_MAPPING.values()],
269
            "label": f"./labelsTr/{case}.nii.gz",
270
        }
271
        for case in training_cases
272
    ]
273
    dataset_json['test'] = [
274
        [f"./imagesTs/{case}_{suffix}.nii.gz" for suffix in MODALITY_MAPPING.values()]
275
        for case in test_cases
276
    ]
277

278
    json_file_path = os.path.join(task_folder, "dataset.json")
279
    with open(json_file_path, 'w', encoding='utf-8') as f:
280
        json.dump(dataset_json, f, indent=4, ensure_ascii=False)
281

282
    print(f"dataset.json 已保存到: {json_file_path}")
283
    print("BraTS2023数据转换完成！标签保持原始值: 0(背景), 1(坏死/非增强), 2(水肿), 3(增强肿瘤)")
284
    return task_folder
285

286

287
def parse_args():
288
    parser = argparse.ArgumentParser(description="BraTS2023 -> nnU-Net 转换（多线程）")
289
    parser.add_argument("--brats-root", default="/root/autodl-tmp/BraTS2023")
290
    parser.add_argument("--nnunet-raw", default="/root/autodl-tmp/data/nnUNet_raw")
291
    parser.add_argument("--task-name", default="Dataset001_BraTS2023")
292
    parser.add_argument("--workers", type=int, default=8,
293
                        help="线程数；机械盘建议 4，NVMe 建议 8~16")
294
    parser.add_argument("--train-ratio", type=float, default=0.8)
295
    parser.add_argument("--verify-labels", action="store_true",
296
                        help="读取每个 seg 的体素值并统计唯一标签（较慢）")
297
    parser.add_argument("--overwrite", action="store_true",
298
                        help="覆盖已存在的目标文件（默认跳过，可断点续跑）")
299
    return parser.parse_args()
300

301

302
if __name__ == "__main__":
303
    args = parse_args()
304
    try:
305
        result = convert_brats2023_to_nnunet(
306
            brats_root=args.brats_root,
307
            nnunet_raw_data_base=args.nnunet_raw,
308
            task_name=args.task_name,
309
            workers=args.workers,
310
            train_ratio=args.train_ratio,
311
            verify_labels=args.verify_labels,
312
            overwrite=args.overwrite,
313
        )
314
        if result:
315
            print(f"\n转换成功完成！输出目录: {result}")
316
            print("可以继续进行nnUNet的预处理和训练步骤。")
317
        else:
318
            print("\n转换失败，请查看错误日志了解详细信息。")
319
    except Exception as e:
320
        print(f"程序执行失败: {str(e)}")
321
        error_file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "error.txt")
322
        with open(error_file_path, 'w', encoding='utf-8') as f:
323
            f.write(f"程序执行失败: {str(e)}\n")
324
        print(f"错误已记录到: {error_file_path}")

BraTS2024#

1
import os
2
import shutil
3
import json
4
import threading
5
from concurrent.futures import ThreadPoolExecutor, as_completed
6

7
import nibabel as nib
8
import numpy as np
9
from collections import OrderedDict
10

11

12
def convert_brats2024_to_nnunet(brats_root, nnunet_raw_data_base, max_workers=8):
13
    """
14
    将BraTS2024数据集转换为nnUNet格式，保持原始标签不变（多线程加速版）
15

16
    Args:
17
        brats_root: BraTS2024原始数据根目录路径
18
        nnunet_raw_data_base: nnUNet原始数据基础目录路径
19
        max_workers: 并行线程数（I/O 密集，建议 4~16，视磁盘性能而定）
20
    """
21

22
    # 设置路径
23
    task_name = "Dataset001_BraTS2024"
24
    task_folder = os.path.join(nnunet_raw_data_base, "nnUNet_raw", task_name)
25

26
    imagesTr_folder = os.path.join(task_folder, "imagesTr")
27
    imagesTs_folder = os.path.join(task_folder, "imagesTs")
28
    labelsTr_folder = os.path.join(task_folder, "labelsTr")
29
    labelsTs_folder = os.path.join(task_folder, "labelsTs")
30

31
    for folder in [imagesTr_folder, imagesTs_folder, labelsTr_folder, labelsTs_folder]:
32
        os.makedirs(folder, exist_ok=True)
33

34
    # 收集所有病例文件夹
35
    all_cases = []
36
    for item in os.listdir(brats_root):
37
        case_path = os.path.join(brats_root, item)
38
        if os.path.isdir(case_path) and item.startswith('BraTS-'):
39
            all_cases.append(item)
40

41
    all_cases.sort()  # 确保顺序一致，训练/测试划分才可复现
42
    total = len(all_cases)
43
    print(f"找到 {total} 个病例，使用 {max_workers} 个线程处理")
44

45
    # 模态映射 - BraTS2024模态
46
    modality_mapping = {
47
        't1n': '0000',  # T1 native (非增强T1)
48
        't1c': '0001',  # T1 contrast enhanced (增强T1)
49
        't2f': '0002',  # T2 FLAIR
50
        't2w': '0003'   # T2 weighted
51
    }
52

53
    # 训练/测试划分点（前80%训练），事先算好，保证并行下划分与原逻辑一致
54
    split_index = total * 0.8
55

56
    # ---------- 纯函数式辅助：只返回结果，绝不修改外部共享变量 ----------
57

58
    def check_file_validity(file_path):
59
        """检查文件是否有效（非空且可读取），返回 (bool, msg)"""
60
        try:
61
            if not os.path.exists(file_path):
62
                return False, "文件不存在"
63
            if os.path.getsize(file_path) == 0:
64
                return False, "文件为空"
65
            nib.load(file_path)  # 只读文件头，验证可解析
66
            return True, "文件有效"
67
        except Exception as e:
68
            return False, f"文件无效: {str(e)}"
69

70
    def copy_image(src_path, dst_path, case_name, modality, errors):
71
        """复制图像文件，出错时把错误追加到本病例的局部 errors 列表"""
72
        try:
73
            if not src_path.endswith('.gz'):
74
                nib.save(nib.load(src_path), dst_path)
75
            else:
76
                shutil.copy2(src_path, dst_path)
77
            return True
78
        except Exception as e:
79
            errors.append(
80
                f"复制图像失败 - 病例: {case_name}, 模态: {modality}, "
81
                f"文件: {src_path}, 错误: {str(e)}"
82
            )
83
            return False
84

85
    def copy_label(src_path, dst_path, case_name, errors, logs):
86
        """复制标签文件并检查标签值分布，错误/日志写入局部列表"""
87
        try:
88
            label_nii = nib.load(src_path)
89
            label_data = label_nii.get_fdata().astype(np.uint8)
90

91
            unique_labels = np.unique(label_data)
92
            logs.append(f"  {os.path.basename(src_path)} 标签值: {unique_labels}")
93

94
            # 验证标签值是否符合BraTS2024规范（0, 1, 2, 3, 4）
95
            valid_labels = {0, 1, 2, 3, 4}
96
            unexpected_labels = set(int(x) for x in unique_labels) - valid_labels
97
            if unexpected_labels:
98
                errors.append(
99
                    f"发现意外标签值 - 病例: {case_name}, 标签值: {unexpected_labels}"
100
                )
101

102
            # 统计各标签的体素数量
103
            label_counts = {int(l): int(np.sum(label_data == l)) for l in unique_labels}
104
            logs.append(f"  标签分布: {label_counts}")
105

106
            # 直接保存，不做任何标签修改
107
            if not src_path.endswith('.gz'):
108
                nib.save(label_nii, dst_path)
109
            else:
110
                shutil.copy2(src_path, dst_path)
111
            return True
112
        except Exception as e:
113
            errors.append(
114
                f"复制标签失败 - 病例: {case_name}, 文件: {src_path}, 错误: {str(e)}"
115
            )
116
            return False
117

118
    # ---------- 单个病例的完整处理逻辑（在工作线程中运行）----------
119

120
    def process_case(i, case_name):
121
        """
122
        处理单个病例，返回一个结果字典：
123
        {'case_name', 'status': 'train'/'test'/'skip', 'errors': [...], 'logs': [...]}
124
        整个函数只读取外部只读变量，产出全部装进返回值，线程安全。
125
        """
126
        errors = []
127
        logs = []
128
        case_folder = os.path.join(brats_root, case_name)
129

130
        if not os.path.exists(case_folder):
131
            errors.append(f"病例文件夹不存在: {case_folder}")
132
            return {'case_name': case_name, 'status': 'skip', 'errors': errors, 'logs': logs}
133

134
        base_name = case_name  # 例如: BraTS-GLI-00000-000
135
        required_files = {
136
            't1n': f"{base_name}-t1n.nii",
137
            't1c': f"{base_name}-t1c.nii",
138
            't2f': f"{base_name}-t2f.nii",
139
            't2w': f"{base_name}-t2w.nii",
140
            'seg': f"{base_name}-seg.nii",
141
        }
142

143
        # 检查文件存在性与有效性，支持 .nii / .nii.gz
144
        files_ok = True
145
        for key, filename in required_files.items():
146
            file_path = os.path.join(case_folder, filename)
147
            gz_file_path = file_path + ".gz"
148

149
            if os.path.exists(file_path):
150
                valid, msg = check_file_validity(file_path)
151
                if not valid:
152
                    errors.append(f"文件无效 - 病例: {case_name}, {filename}: {msg}")
153
                    files_ok = False
154
            elif os.path.exists(gz_file_path):
155
                required_files[key] = filename + ".gz"
156
                valid, msg = check_file_validity(gz_file_path)
157
                if not valid:
158
                    errors.append(f"文件无效 - 病例: {case_name}, {filename}.gz: {msg}")
159
                    files_ok = False
160
            else:
161
                errors.append(f"文件缺失 - 病例: {case_name}, 文件: {filename} 或 {filename}.gz")
162
                files_ok = False
163

164
        if not files_ok:
165
            return {'case_name': case_name, 'status': 'skip', 'errors': errors, 'logs': logs}
166

167
        # 决定训练还是测试（前80%训练）
168
        is_training = i < split_index
169
        img_dst_folder = imagesTr_folder if is_training else imagesTs_folder
170
        lbl_dst_folder = labelsTr_folder if is_training else labelsTs_folder
171

172
        case_success = True
173

174
        # 复制四个模态图像
175
        for modality, suffix in modality_mapping.items():
176
            src_file = os.path.join(case_folder, required_files[modality])
177
            dst_file = os.path.join(img_dst_folder, f"{case_name}_{suffix}.nii.gz")
178
            if not copy_image(src_file, dst_file, case_name, modality, errors):
179
                case_success = False
180

181
        # 复制分割标签
182
        src_seg = os.path.join(case_folder, required_files['seg'])
183
        dst_seg = os.path.join(lbl_dst_folder, f"{case_name}.nii.gz")
184
        if not copy_label(src_seg, dst_seg, case_name, errors, logs):
185
            case_success = False
186

187
        if case_success:
188
            status = 'train' if is_training else 'test'
189
        else:
190
            status = 'skip'
191
        return {'case_name': case_name, 'status': status, 'errors': errors, 'logs': logs}
192

193
    # ---------- 并行调度 ----------
194

195
    training_cases = []
196
    test_cases = []
197
    skipped_cases = []
198
    errors = []
199

200
    done = 0
201
    progress_lock = threading.Lock()
202

203
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
204
        futures = {
205
            executor.submit(process_case, i, case_name): case_name
206
            for i, case_name in enumerate(all_cases)
207
        }
208

209
        for future in as_completed(futures):
210
            case_name = futures[future]
211
            try:
212
                result = future.result()
213
            except Exception as e:
214
                # 兜底：worker 内部未捕获的异常
215
                errors.append(f"病例 {case_name} 处理时发生未捕获异常: {str(e)}")
216
                skipped_cases.append(case_name)
217
                result = None
218

219
            if result is not None:
220
                errors.extend(result['errors'])
221
                status = result['status']
222
                if status == 'train':
223
                    training_cases.append(result['case_name'])
224
                elif status == 'test':
225
                    test_cases.append(result['case_name'])
226
                else:
227
                    skipped_cases.append(result['case_name'])
228

229
            # 线程安全地更新并打印进度
230
            with progress_lock:
231
                done += 1
232
                print(f"[{done}/{total}] {case_name} -> "
233
                      f"{result['status'] if result else 'skip'}")
234

235
    # 结果列表排序，保证输出稳定（并行完成顺序是乱的）
236
    training_cases.sort()
237
    test_cases.sort()
238
    skipped_cases.sort()
239

240
    print(f"\n处理完成: {len(training_cases)} 个训练病例, {len(test_cases)} 个测试病例")
241
    print(f"跳过的病例数量: {len(skipped_cases)}")
242

243
    # ---------- 写入错误日志 ----------
244
    error_file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "error.txt")
245
    with open(error_file_path, 'w', encoding='utf-8') as f:
246
        f.write("BraTS2024数据转换错误报告\n")
247
        f.write("=" * 80 + "\n\n")
248
        f.write(f"总计处理: {total} 个病例\n")
249
        f.write(f"成功处理: {len(training_cases) + len(test_cases)} 个病例\n")
250
        f.write(f"跳过病例: {len(skipped_cases)} 个病例\n")
251
        f.write(f"错误数量: {len(errors)} 个错误\n\n")
252

253
        if skipped_cases:
254
            f.write("跳过的病例列表:\n")
255
            for case in skipped_cases:
256
                f.write(f"  - {case}\n")
257
            f.write("\n")
258

259
        if errors:
260
            f.write("详细错误信息:\n")
261
            for idx, error in enumerate(errors, 1):
262
                f.write(f"{idx}. {error}\n")
263
        else:
264
            f.write("没有发现错误。\n")
265

266
    print(f"错误日志已保存到: {error_file_path}")
267

268
    # ---------- 生成 dataset.json（nnU-Net v2 格式）----------
269
    if len(training_cases) + len(test_cases) > 0:
270
        dataset_json = OrderedDict()
271

272
        # v2: channel_names 取代 v1 的 modality，键为通道索引字符串
273
        dataset_json['channel_names'] = {
274
            "0": "T1n",
275
            "1": "T1c",
276
            "2": "T2f",
277
            "3": "T2w",
278
        }
279

280
        # v2: labels 是 名称 -> 整数值（与 v1 正好相反）
281
        dataset_json['labels'] = {
282
            "background": 0,
283
            "NETC": 1,   # Non-Enhancing Tumor Core
284
            "SNFH": 2,   # Surrounding Non-enhancing FLAIR Hyperintensity
285
            "ET": 3,     # Enhancing Tumor
286
            "RC": 4,     # Resection Cavity (BraTS2024 新增)
287
        }
288

289
        # v2 必需字段
290
        dataset_json['numTraining'] = len(training_cases)
291
        dataset_json['file_ending'] = ".nii.gz"
292

293
        # 以下为可选元信息，nnU-Net 不使用但不会报错，仅方便追溯
294
        dataset_json['name'] = "BraTS2024"
295
        dataset_json['description'] = "Brain Tumor Segmentation Challenge 2024"
296
        dataset_json['reference'] = "https://www.synapse.org/#!Synapse:syn53708249"
297
        dataset_json['licence'] = "see BraTS2024 website"
298
        dataset_json['release'] = "1.0"
299

300
        json_file_path = os.path.join(task_folder, "dataset.json")
301
        with open(json_file_path, 'w') as f:
302
            json.dump(dataset_json, f, indent=4)
303

304
        print(f"dataset.json 已保存到: {json_file_path}")
305
        print("BraTS2024数据转换完成！标签保持原始值不变：")
306
        print("  0 - background (背景)")
307
        print("  1 - NETC - 非增强肿瘤核心")
308
        print("  2 - SNFH - 周围非增强FLAIR高信号")
309
        print("  3 - ET - 增强肿瘤")
310
        print("  4 - RC - 切除腔")
311
        return task_folder
312
    else:
313
        print("警告: 没有成功处理任何病例，未生成dataset.json文件")
314
        return None
315

316

317
# 使用示例
318
if __name__ == "__main__":
319
    brats_root = "BraTS2024"       # 您的BraTS2024数据根目录
320
    nnunet_raw_data_base = "data"  # nnUNet原始数据基础目录
321

322
    try:
323
        result = convert_brats2024_to_nnunet(brats_root, nnunet_raw_data_base, max_workers=8)
324
        if result:
325
            print(f"\n转换成功完成！输出目录: {result}")
326
            print("可以继续进行nnUNet的预处理和训练步骤。")
327
            print("\n注意：BraTS2024引入了新的切除腔（RC）标签（标签值4）。")
328
        else:
329
            print("\n转换失败，请查看错误日志了解详细信息。")
330
    except Exception as e:
331
        print(f"程序执行失败: {str(e)}")
332
        error_file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "error.txt")
333
        with open(error_file_path, 'w', encoding='utf-8') as f:
334
            f.write(f"程序执行失败: {str(e)}\n")
335
        print(f"错误已记录到: {error_file_path}")

nnUNetv2数据集转换

BraTS数据集转换成MSD数据集#

BraTS2019:#

BraTS2023#

BraTS2024#

目录