nnUNet V2代码——生成dataset_fingerprint.json

原创已于 2025-02-23 12:00:50 修改 · 1.6k 阅读

13 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能 #计算机视觉 #深度学习 #机器学习 #卷积神经网络 #nnU-Net V2 #nnunet

于 2024-08-14 10:17:17 首次发布

阅读nnUNet V2代码专栏收录该内容

15 篇文章

订阅专栏

前文请见nnUNetv2_plan_and_preprocess命令

阅读nnUNet\nnunetv2\experiment_planning\dataset_fingerprint\fingerprint_extractor.py

文件内只有一个DatasetFingerprintExtractor类，阅读它

在DatasetFingerprintExtractor类内涉及的其他函数都在文章后半部分说明

一. init函数

1. 参数

dataset_name_or_id：数据集名称或ID
num_processes：进程数
verbose：是否打印详细信息

2. 过程

定义一些变量，部分如下：

dataset_name：数据集名称，例如Dataset001_QQQQ
dataset_json：数据集相关信息，是用户创建的，例如：Dataset001_QQQQ/dataset_json
dataset：字典类型，{“样例1” : { “images” : {原始图像文件路径列表}, “label” : {分割掩码文件路径列表} }}
num_foreground_voxels_for_intensitystats：数据集一共会采样的前景像素数量，用于统计，后续有介绍

二.run函数

1. 参数

无

2. 过程

获取nnUNet_preprocessed文件夹路径，再创建nnUNet_preprocessed文件夹，最后获取dataset_fingerprint.json文件路径。代码较为清晰，不做粘贴

由上述dataset_fingerprint.json文件路径判断其是否已经存在 且不选择覆写，则载入现有文件并返回。代码较为清晰，不做粘贴

如果fingerprint文件不存在 或选择覆写，则：

首先获取本次训练使用数据集的读取和写入器（由determine_reader_writer_from_dataset_json函数函数获取，实质是查表，见本文最后）

接下来计算每个样例的前景体素****采样次数，再利用多进程运行analyze_case函数，快速处理数据集多个样例。多进程代码不做粘贴。

从此处继续，results列表存储了每一例数据调用analyze_case函数返回的结果：

results = [i.get()[0] for i in r]

从results列表内提取相关变量，依据这些变量，计算数据集的均值、最值等信息，存入fingerprint变量并写入dataset_fingerprint.json文件

三.analyze_case函数

1. 参数

image_files：image
segmentation_file：seg
reader_writer_class：样例文件读取和写入器
num_samples：前景体素采样次数

2. 过程

本函数统计图像大小、体素间距等信息

读取image和seg文件，按照image非零区域对image和seg裁剪，获得以下三个变量：

data_cropped, seg_cropped, bbox = crop_to_nonzero(images, segmentation)

接下来使用collect_foreground_intensities函数收集裁剪后image的统计信息

这些信息包括：

mean	median	min	max	percentile_99_5	percentile_00_5
均值	中位数	最小值	最大值	第99.5百分位数	第0.5百分位数

foreground_intensities_per_channel, foreground_intensity_stats_per_channel = \
            DatasetFingerprintExtractor.collect_foreground_intensities(seg_cropped, data_cropped,
                                                                       num_samples=num_samples)

获取image体素间距（spacing）、裁剪前后的image大小（shape_before_crop和shape_after_crop）以及裁剪比例（relative_size_after_cropping）：

spacing = properties_images['spacing']

shape_before_crop = images.shape[1:]
shape_after_crop = data_cropped.shape[1:]
relative_size_after_cropping = np.prod(shape_after_crop) / np.prod(shape_before_crop)

最后返回上述相关信息：

四.collect_foreground_intensities函数

1. 参数

segmentation：seg
images：image
seed：随机采样种子
num_samples：采样数

2. 过程

此函数统计image中seg前景区域内均值、最值等信息

首先进行断言和创建随机数生成器，随机数生成器用于随机采样，代码较为清晰，不做粘贴。

初始化intensities_per_channel和intensity_statistics_per_channel列表，二者分别用于记录随机采取的image前景像素值和image所有采样像素值的均值、最值等信息；去除seg中的空维度，并获取seg中非背景区域，存入foreground_mask变量；定义用于统计百分位数的数组，部分代码没有粘贴：

foreground_mask = segmentation[0] > 0

和去除掩码图的空维度类似，对图像的第一个维度遍历（我使用的数据集是ACDC数据集，此处的image第一个维度和seg第一个维度一样，都是空维度）。for循环开始：

images[i]的shape是channel * height * width；获取images[i]中在seg前景区域内的像素；再利用rs随机采样器对这些像素随机、有放回的采样；

intensities_per_channel.append(
        rs.choice(foreground_pixels, num_samples, replace=True) if num_fg > 0 else [])

由于有些样例的seg并没有前景区域（全是背景，没有前景），所以作者默认均值最值等信息是nan，如果含有前景区域则统计信息并存入intensity_statistics_per_channel变量；

mean, median, mini, maxi, percentile_99_5, percentile_00_5 = np.nan, np.nan, np.nan, np.nan, np.nan, np.nan

其余代码较为清晰，不做粘贴

for循环结束。

最后返回上述计算的数据

五.涉及的函数

1. get_filenames_of_train_images_and_targets函数

a. 参数

raw_dataset_folder：nnUNet_raw下的数据集文件夹，例如nnunet_raw/Dataset001_QQQQ
dataset_json：用户创建的dataset_json

b. 过程

首先是确保成功读取dataset.json文件

如果dataset.json文件中存在dataset键值，则规范化用户定义的文件路径

否则根据raw_dataset_folder文件夹获取正确的数据集内各个样例的文件路径。

最后将数据集各样例标识符（区分样例和样例的）、image文件路径、seg文件路径打包成字典并返回

if 'dataset' in dataset_json.keys():
    dataset = dataset_json['dataset']
    for k in dataset.keys():
        dataset[k]['label'] = os.path.abspath(join(raw_dataset_folder, dataset[k]['label'])) if not os.path.isabs(dataset[k]['label']) else dataset[k]['label']
        dataset[k]['images'] = [os.path.abspath(join(raw_dataset_folder, i)) if not os.path.isabs(i) else i for i in dataset[k]['images']]
else:
    identifiers = get_identifiers_from_splitted_dataset_folder(join(raw_dataset_folder, 'imagesTr'), dataset_json['file_ending'])
    images = create_lists_from_splitted_dataset_folder(join(raw_dataset_folder, 'imagesTr'), dataset_json['file_ending'], identifiers)
    segs = [join(raw_dataset_folder, 'labelsTr', i + dataset_json['file_ending']) for i in identifiers]
    dataset = {i: {'images': im, 'label': se} for i, im, se in zip(identifiers, images, segs)}

2. get_identifiers_from_splitted_dataset_folder函数

a. 参数

folder：待读取的文件夹
file_ending：样例文件的拓展名

b. 过程

读取文件夹下的所有文件

去除文件名的拓展位和**_XXXX**，例如：QQQQQ001_0000.nii.gz --> QQQQQ001，这样就获取了每个样例的唯一标识名

去重后返回标识名列表

def get_identifiers_from_splitted_dataset_folder(folder: str, file_ending: str):
    files = subfiles(folder, suffix=file_ending, join=False)
    # all files have a 4 digit channel index (_XXXX)
    crop = len(file_ending) + 5
    files = [i[:-crop] for i in files]
    # only unique image ids
    files = np.unique(files)
    return files

3. create_lists_from_splitted_dataset_folder函数

a. 参数

无

b. 过程

确保identifiers是folder下所有文件的标识名列表

获取folder下所有文件名称，存入files变量

遍历所有标识名，如果符合“标识名 _XXXX.拓展名”的格式，就将其文件路径加入list_of_lists列表

遍历结束返回list_of_lists列表

def create_lists_from_splitted_dataset_folder(folder: str, file_ending: str, identifiers: List[str] = None) -> List[
    List[str]]:
    """
    does not rely on dataset.json
    """
    if identifiers is None:
        identifiers = get_identifiers_from_splitted_dataset_folder(folder, file_ending)
    files = subfiles(folder, suffix=file_ending, join=False, sort=True)
    list_of_lists = []
    for f in identifiers:
        p = re.compile(re.escape(f) + r"_\d\d\d\d" + re.escape(file_ending))
        list_of_lists.append([join(folder, i) for i in files if p.fullmatch(i)])
    return list_of_lists

4. determine_reader_writer_from_dataset_json函数

a. 参数

dataset_json_content：内容是dataset.json
example_file：数据集的一个数据
allow_nonmatching_filename：是否允许不匹配的文件名
verbose：是否输出详细信息

b. 过程

判断“overwrite_image_reader_writer”是否在dataset.json文件内设置，就是用户是否指定了文件读写器

如果设置且不为None，则按名称寻找

否则调用determine_reader_writer_from_file_ending函数，按照文件拓展名自动寻找读取器和写入器（实际是查表）

def determine_reader_writer_from_dataset_json(dataset_json_content: dict, example_file: str = None,
                                              allow_nonmatching_filename: bool = False, verbose: bool = True
                                              ) -> Type[BaseReaderWriter]:
    if 'overwrite_image_reader_writer' in dataset_json_content.keys() and \
            dataset_json_content['overwrite_image_reader_writer'] != 'None':
        ioclass_name = dataset_json_content['overwrite_image_reader_writer']
        # trying to find that class in the nnunetv2.imageio module
        try:
            ret = recursive_find_reader_writer_by_name(ioclass_name)
            if verbose: print(f'Using {ret} reader/writer')
            return ret
        except RuntimeError:
            # 略
    return determine_reader_writer_from_file_ending(dataset_json_content['file_ending'], example_file,
                                                    allow_nonmatching_filename, verbose)