brats2019_clean/dataclean_BRATS_2019.py

#coding:utf-8
'''
write by ygq
create on 2025-08-03
update BRATS_2019

BRATS2019 是一个大规模、多模态、标注精良的脑胶质瘤 MRI 数据集，主要用于开发和评估脑肿瘤自动分割算法以及基于 MRI 的生存预测模型。
它包含 335 例带标注的训练数据和大量未标注的验证/测试数据。其多模态特性（T1, T1Gd, T2, T2-FLAIR）和精细的肿瘤子区域标注（整个肿瘤、肿瘤核心、增强肿瘤）
训练集： 包含 335 例 患者的完整多模态 MRI 扫描数据及其对应的专家手动分割标注（Ground Truth）

每个病例包含四种不同序列的 3D MRI 扫描（均已进行预处理，如配准、重采样到 1mm³ 各向同性、颅骨剥离）：
    Native (T1): 标准的 T1 加权成像，显示解剖结构。
    Post-contrast T1-weighted (T1Gd/T1ce): 注射钆造影剂后的 T1 加权成像。造影剂会渗漏过被破坏的血脑屏障，在肿瘤活跃区域（如坏死核心的边缘）呈现强化。
    T2-weighted (T2): T2 加权成像，对水肿和囊性/坏死区域非常敏感，呈现高信号。
    T2 Fluid Attenuated Inversion Recovery (T2-FLAIR): 抑制了脑脊液信号的 T2 加权成像。特别擅长显示肿瘤周围的水肿区域（通常也包含浸润性肿瘤细胞），呈现高信号。

训练集提供了由专业医师手动精细勾画的肿瘤区域分割标注。
标注定义了三个相互嵌套或重叠的子区域（反映了肿瘤的不同生物学特性）：
    坏疽性和非增强肿瘤核心： 包括坏死区域（在 T1Gd 上无强化）和活跃肿瘤的非增强部分（在 T2-FLAIR 上高信号，但在 T1Gd 上不强化）。标签值 = 1。
    瘤周水肿： 肿瘤周围的水肿区域（在 T2 和 T2-FLAIR 上呈高信号）。标签值 = 2。
    增强肿瘤： 在 T1Gd 上呈现强化的区域（通常代表高度血管化的活跃肿瘤组织）。标签值 = 4。
    整个肿瘤区域 由这三个区域组合而成（标签值 1+2+4）。
    肿瘤核心区域 由坏疽性和非增强肿瘤核心 + 增强肿瘤组成（标签值 1+4）

    对于训练集和验证集中的 HGG 患者，提供了患者的总生存期信息（从初次扫描到死亡或最后一次随访的时间）


根据沟通参考MSD中的BRATS的结构：
    1.将多个分开的模态合并，构建第四个维度的数组，分别按照FLAIR,T1,T1CE,T2顺序存放；
    2.生存期信息也需要相应补充道HGG的数据集中


meta_info:[保留Grade,BraTS_2019_subject_ID]
    Grade,BraTS_2017_subject_ID,BraTS_2018_subject_ID,TCGA_TCIA_subject_ID,BraTS_2019_subject_ID
    HGG,Brats17_CBICA_AAB_1,Brats18_CBICA_AAB_1,NA,BraTS19_CBICA_AAB_1
    LGG,Brats17_TCIA_451_1,Brats18_TCIA09_451_1,TCGA-CS-4942,BraTS19_TCIA09_451_1
    
survival_info:
    BraTS19ID,Age,Survival,ResectionStatus
    BraTS19_CBICA_AAB_1,60.4630137,289,GTR
    BraTS19_CBICA_AAG_1,52.2630137,616,GTR
'''
import os
import glob
import pandas as pd
import SimpleITK as sitk
import argparse
import json
from tqdm import tqdm
from util import meta_data
import util
import numpy as np
# from bert_helper import *

import shutil
##dataset_meta
meta_id_name='BraTS_2019_subject_ID'
meta_grade_name='Grade'

##HGG_survival_info
survival_id_name='BraTS19ID'
meta_age_name='Age'
meta_survival_name='Survival'
meta_status_name='ResectionStatus'


TASK_VALUE="segmentation"
CLAMP_RANGE_CT = [-300,300]
CLAMP_RANGE_MRI = None # MRI images threshold placeholder TBC...
TARGET_VOXEL_SPACING=None

##参考MSD的sub_modality描述信息
SUB_MODALITY=["FLAIR","T1w","t1gd","T2w"]
##文件名对应的排序顺序
SERIES_ORDER=["flair","t1","t1ce","t2"]

LABEL_DICT={
    "0":"backgroud",
    "1":"non-enhancing tumor",
    "2":"edema",
    "4":"enhancing tumour"
}
# def find_metadata_files(path):
#     # for Cancer Image Archive (TCIA) dataset
#     search_pattern = os.path.join(path, '**', 'metadata.csv')
#     return glob.glob(search_pattern, recursive=True)

def find_metadata_files(path):
    # for Cancer Image Archive (TCIA) dataset
    search_pattern = os.path.join(path, '*.csv')
    return glob.glob(search_pattern, recursive=True)
##added by yanguoqing on 20250527
def find_image_dirs(path):
    return os.listdir(path)

##modify by yanguoqing on 20250527
def load_dicom_images(folder_path):
    reader = sitk.ImageSeriesReader()
    dicom_names = reader.GetGDCMSeriesFileNames(folder_path)
    reader.SetFileNames(dicom_names)
    image = reader.Execute()
    return dicom_names,image

##added by yanguoqing on 20250527
def load_dicom_tag(imgs):
    reader = sitk.ImageFileReader()
    # dicom_names = reader.GetGDCMSeriesFileNames(folder_path)
    reader.SetFileName(imgs)
    reader.ReadImageInformation()  # 仅读取元信息，不加载像素数据
    # metadata_keys = reader.GetMetaDataKeys()
    tag=reader.Execute()
    return tag

def load_nrrd(fp):
    return sitk.ReadImage(fp)

##modify by yanguoqing on 20250805
def load_brtas_images(series_files):
    '''
    每个病例包含四种不同序列的 3D MRI 扫描（均已进行预处理，如配准、重采样到 1mm³ 各向同性、颅骨剥离）
    将多个分开的模态合并，构建第四个维度的数组，分别按照FLAIR,T1,T1CE,T2顺序存放
    '''
    reader = sitk.ImageSeriesReader()
    reader.SetFileNames(series_files)
    image = reader.Execute()
    return image

def save_nifti(image, output_path, folder_path):
    # Set metadata in the NIfTI file's header
    output_dirpath = os.path.dirname(output_path)
    if not os.path.exists(output_dirpath):
        print(f"Creating directory {output_dirpath}")
        os.makedirs(output_dirpath)
    # Set metadata in the NIfTI file's header
    image.SetMetaData("FolderPath", folder_path)
    sitk.WriteImage(image, output_path)

##modify by yanguoqing on 20250527
def convert_windows_to_linux_path(windows_path):
    # Replace backslashes with forward slashes and remove the drive letter
    # Some meta files have windows paths, but the data is stored on a linux server
    linux_path = windows_path.replace('\\', '/')
    if ':' in linux_path:
        linux_path = linux_path.split(':', 1)[1]
    return linux_path

def main(target_path, output_dir):
    metadata_files = find_metadata_files(target_path)
    pid_dirs=find_image_dirs(target_path)
    pid_dirs=["HGG","LGG"]
    failed_files = []
    if not os.path.isdir(output_dir):
        os.makedirs(output_dir)
    json_output_path = os.path.join(output_dir, 'nifti_mappings.json')
    failed_files_path = os.path.join(output_dir, 'failed_files.json')
    meta = meta_data()
    
    # Initialize the JSON file
    if not os.path.exists(json_output_path):
        with open(json_output_path, 'w') as json_file:
            json.dump({}, json_file)
    meta_file=os.path.join(target_path,'name_mapping.csv')
    survival_file=os.path.join(target_path,'survival_data.csv')
    if os.path.isfile(meta_file):
        mf_flag=True
        df_meta=pd.read_csv(meta_file,sep=',')
    else:
        mf_flag=False

    if os.path.isfile(survival_file):
        sf_flag=True
        df_survial=pd.read_csv(survival_file,sep=',')
    else:
        sf_flag=False

    if pid_dirs:
        for pid_dir in tqdm(pid_dirs, desc="Processing pid dirs"):
            if not os.path.isdir(os.path.join(target_path,pid_dir)):
                continue
            ##HGG_FLAG
            if pid_dir =="HGG":
                tr_flag=True
            else:
                tr_flag=False
            # label_flag=False
            
            ##遍历所有目录下的HGG/LGG的病例数据（影像+标注seg）
            image_dirs=find_image_dirs(os.path.join(target_path,pid_dir))
            
            for data_dir in tqdm(image_dirs, desc="Processing images files"):
                full_path=os.path.join(target_path,pid_dir,data_dir)
                
                data_info_row=df_meta[df_meta[meta_id_name]==data_dir]
                
                if data_info_row.shape[0]>0:
                    data_info_row=data_info_row.reset_index()
                    #print(data_info_row[meta_id_name])
                    meta_image_id=data_info_row[meta_id_name][0]
                    meta_image_grade=data_info_row[meta_grade_name][0]

                else:
                    meta_image_id=data_dir
                    meta_image_grade=''
                
                if tr_flag:
                    survival_file_row=df_survial[df_survial[survival_id_name]==data_dir]
                    if data_info_row.shape[0]>0:
                        survival_file_row=survival_file_row.reset_index()
                        #print(data_info_row[meta_id_name])
                        meta_image_age=survival_file_row[meta_age_name][0]
                        meta_image_survival=survival_file_row[meta_survival_name][0]
                        meta_image_status=survival_file_row[meta_status_name][0]

                    else:
                        meta_image_age=''
                        meta_image_survival=''
                        meta_image_status=''
                else:
                    meta_image_age=''
                    meta_image_survival=''
                    meta_image_status=''
                

                try:
                    ##读取MRI四组文件，按照flair,t1,t1ce,t2的顺序叠加，对于seg先剔除不参与

                    
                    series_files=[os.path.join(full_path,"%s_%s.nii"%(data_dir,sm))for sm in SERIES_ORDER]
                    ##判断是否每个sub_modality文件存在
                    series_flag=[os.path.isfile(os.path.join(full_path,"%s_%s.nii"%(data_dir,sm)))for sm in SERIES_ORDER]
                    series_files=[series_files[index] for index, value in enumerate(series_flag) if  value]
                    sub_modality=[SUB_MODALITY[index] for index, value in enumerate(series_flag) if  value]
                    if len(series_files)>0:
                        ##存在有效的MRI影像数据进行后续处理
                        sitk_img_original=load_brtas_images(series_files)
                        
                    else:
                        print("病例数据%s为空"%data_dir)
                        continue
                
                    
                    original_spacing = list(sitk_img_original.GetSpacing())
                    original_size = list(sitk_img_original.GetSize())

                    modality="MRI"
                    study='BRATS_2019'##Dataset_name
                    CIA_other_info = {
                        'metadata_file':''
                
                    }
                    CIA_other_info['split'] = "train"
                    if mf_flag:
                        CIA_other_info['metadata_file']=meta_file
                    

                    ##
                    CIA_other_info['Image_id']=meta_image_id
                    CIA_other_info['Grade']=meta_image_grade
                    CIA_other_info['Age']=str(meta_image_age)
                    CIA_other_info['Survival']=str(meta_image_survival)
                    CIA_other_info['ResectionStatus']=meta_image_status


                    meta.add_keyvalue('Spacing_mm',1.0)
                    meta.add_keyvalue('OriImg_path',",".join(series_files))
                    meta.add_keyvalue('Size',original_size)  # 这里用处理后的size -- YH Jachin
                    meta.add_keyvalue('Modality',modality)
                    meta.add_keyvalue('Dataset_name',study)
                    meta.add_keyvalue('ROI','head')

                    sub_modality_dict={}
                    for idx,value in enumerate(series_flag):
                        if value:
                            sub_modality_dict[str(idx)]=SUB_MODALITY[idx]

                    meta.add_keyvalue('Sub_modality',sub_modality_dict)

                    meta.add_keyvalue('Label_Dict',LABEL_DICT)

                    output_image_file = os.path.join(output_dir,data_dir, f"{data_dir}.nii.gz")
                    # output_path=convert_windows_to_linux_path(output_path)
                    ##
                    save_nifti(sitk_img_original, output_image_file, full_path)
                    print(f"Saved NIfTI file to {output_image_file}")
                    ##Label processing

                    label_path_dict={}
                    full_label_file=os.path.join(full_path,"%s_seg.nii"%(data_dir))

                    process_label_path=os.path.join(output_dir,data_dir,'segmentation')

                    processed_lbl_full_path=os.path.join(process_label_path, f"{data_dir}.nii.gz")

                    if not os.path.isdir(process_label_path):
                        os.makedirs(process_label_path,exist_ok=True)

                    if not os.path.isfile(full_label_file):
                        pass
                        label_flag=False
                    else:
                        sitk_lbl_original = util.load_nifti(full_label_file)
                        util.save_nifti(sitk_lbl_original, processed_lbl_full_path, full_label_file) # Save original
                        print(f"Saved Segemention NIfTI file to {processed_lbl_full_path}")

                        label_path_dict['brain'] = processed_lbl_full_path
                        label_flag=True
                    
                    if label_flag:
                        meta.add_keyvalue('Task',TASK_VALUE)
                        meta.add_keyvalue('Label_path',{TASK_VALUE:label_path_dict})


               
                #     try:
                #         assert sitk_img_processed.GetSize() == sitk_lbl_processed.GetSize()
                #     except Exception as e:
                #         failed_files.append(full_path_label)
                #         continue
                    print(sitk_img_original.GetSize(),sitk_lbl_original.GetSize())

                except Exception as e:
                    print(e)
                    failed_files.append(data_dir)
                    print(f"Failed to load BRATS images from {data_dir}")
                    continue

                
                
                meta.add_extra_keyvalue('Metadata',CIA_other_info)


                # Write the mapping to the JSON file on the fly
                with open(json_output_path, 'r+') as json_file:
                    existing_mappings = json.load(json_file)
                    existing_mappings[output_image_file] = meta.get_meta_data()
                    json_file.seek(0)
                    # print(existing_mappings)
                    json.dump(existing_mappings, json_file, indent=4)
                    json_file.truncate()
    # else:
    #     print("No metadata.csv files found.")
    
    with open(failed_files_path, "w") as json_file:
        json.dump(failed_files, json_file)
        
    print(f"The list has been written to {failed_files_path}")
    print(f"Saved NIfTI mappings to {json_output_path}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Process DICOM files and save as NIfTI.")
    parser.add_argument("--target_path", type=str, help="Path to the target directory containing metadata files.", default="/home/data/Github/data/data_gen_def/DATASETS/BRATS/BRATS2019/MICCAI_BraTS_2019_Data_Training/")
    parser.add_argument("--output_dir", type=str, help="Directory to save the NIfTI files.", default="/home/data/Github/data/data_gen_def/DATASETS_processed/BRATS/BRATS2019")
    args = parser.parse_args()
    print(args.target_path, args.output_dir)
    main(args.target_path, args.output_dir)