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 shutil
from 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_holes
def 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] = 1
seg[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] = 0
io.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=download
source = "/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 segmentation
valid_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 set
valid_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,
)