nnunetv2系列：自定义2D实例分割数据集转换
这里主要参考官方源文件nnUNet/nnunetv2/dataset_conversion/Dataset120_RoadSegmentation.py，注释了一些不必要的操作。数据集下载链接: massachusetts-roads-dataset
重要提示:
nnU-Net只能用于使用无损(或无)压缩的文件格式!因为文件格式是为整个数据集定义的(而不是单独为图像和分割定义的，这可能是将来的任务)，我们必须确保没有破坏分割映射的压缩工件。所以没有。jpg之类的!
支持的2D数据集文件类型包括.png、.bmp、.tif。
原数据集目录结构
这里展示massachusetts-roads-dataset数据集的目录结构。
这里的testing目录在训练过程中并不会使用，默认是在training目录中划分训练集和验证集。
./datasets/road_segmentation_ideal/
├── testing/
│ ├── input/
│ │ ├── img-10.png
│ │ ├── img-11.png
│ │ ├── ...
│ └── output/
│ ├── img-10.png
│ ├── img-11.png
│ │ ├── ...
└── training/├── input/│ ├── img-1000.png│ ├── img-1001.png
│ │ ├── ...└── output├── img-1000.png├── img-1002.png
│ │ ├── ...转换后数据集目录结构
nnUNet_raw/Dataset120_RoadSegmentation
├── dataset.json
├── imagesTr
│ ├── img-2_0000.png
│ ├── img-7_0000.png
│ ├── ...
├── imagesTs # optional
│ ├── img-1_0000.png
│ ├── img-2_0000.png
│ ├── ...
└── labelsTr
| ├── img-2.png
| ├── img-7.png
| ├── ...
└── labelsTs
| ├── img-1.png
| ├── img-2.png
| ├── ...
转换代码示例
这里展示的是包括背景在内的三类实例分割数据集转换代码。
import multiprocessing
import shutilfrom batchgenerators.utilities.file_and_folder_operations import (join,maybe_mkdir_p,subfiles,
)from nnunetv2.dataset_conversion.generate_dataset_json import (generate_dataset_json,
)
from nnunetv2.paths import nnUNet_raw
from skimage import io
# from acvl_utils.morphology.morphology_helper import generic_filter_components
# from scipy.ndimage import binary_fill_holesdef load_and_covnert_case(input_image: str,input_seg: str,output_image: str,output_seg: str,min_component_size: int = 50,
):seg = io.imread(input_seg)seg[seg == 128] = 1seg[seg == 255] = 2# image = io.imread(input_image)# image = image.sum(2)# mask = image == (3 * 255)# # the dataset has large white areas in which road segmentations can exist but no image information is available.# # Remove the road label in these areas# mask = generic_filter_components(# mask,# filter_fn=lambda ids, sizes: [# i for j, i in enumerate(ids) if sizes[j] > min_component_size# ],# )# mask = binary_fill_holes(mask)# seg[mask] = 0io.imsave(output_seg, seg, check_contrast=False)shutil.copy(input_image, output_image)if __name__ == "__main__":# extracted archive from https://www.kaggle.com/datasets/insaff/massachusetts-roads-dataset?resource=downloadsource = "/home/bio/family/segmenation/nnUNet/datasets/eye_sclera_iris_segmentation"dataset_name = "Dataset500_ScleraIrisSegmentation"imagestr = join(nnUNet_raw, dataset_name, "imagesTr")imagests = join(nnUNet_raw, dataset_name, "imagesTs")labelstr = join(nnUNet_raw, dataset_name, "labelsTr")labelsts = join(nnUNet_raw, dataset_name, "labelsTs")maybe_mkdir_p(imagestr)maybe_mkdir_p(imagests)maybe_mkdir_p(labelstr)maybe_mkdir_p(labelsts)train_source = join(source, "training")test_source = join(source, "testing")with multiprocessing.get_context("spawn").Pool(8) as p:# not all training images have a segmentationvalid_ids = subfiles(join(train_source, "output"), join=False, suffix="png")num_train = len(valid_ids)r = []for v in valid_ids:r.append(p.starmap_async(load_and_covnert_case,((join(train_source, "input", v),join(train_source, "output", v),join(imagestr, v[:-4] + "_0000.png"),join(labelstr, v),50,),),))# test setvalid_ids = subfiles(join(test_source, "output"), join=False, suffix="png")for v in valid_ids:r.append(p.starmap_async(load_and_covnert_case,((join(test_source, "input", v),join(test_source, "output", v),join(imagests, v[:-4] + "_0000.png"),join(labelsts, v),50,),),))_ = [i.get() for i in r]generate_dataset_json(join(nnUNet_raw, dataset_name),{0: "R", 1: "G", 2: "B"},{"background": 0, "iris": 1, "sclera": 2},num_train,".png",dataset_name=dataset_name,)
)
