MIC-DKFZ / nnUNet

Apache License 2.0
5.95k stars 1.77k forks source link

Merging segmentations #2570

Closed Mai-CS closed 2 weeks ago

Mai-CS commented 3 weeks ago

Thanks for the great work. I am trying to train nnUNet on a different dataset where each CT scan (.nrrd) has 30 labels (representing 30 organs), so overall, I have 30 .nrrd files. However, nnUNet works when we have one combined file for all segmentations. How can that be handled?

fitzjalen commented 2 weeks ago

You have to combine all this labels into 1 file (I use .nii.gz). The code structure will be depend on the structure of your directory but in basic the good way is to:

  1. Create dictionary with name of the segmentation (same as in the nrrd file) and corresponding label

`label_dict = { 'liver': 1, 'lesion': 2, 'vessel': 3

Add more labels as needed

}`

  1. Write the algorithm that will run throw your directory and for every case create combined mask

`def combine_labels(input_dir, output_dir):

Ensure output directory exists

os.makedirs(output_dir, exist_ok=True)

for case in os.listdir(input_dir):
    case_path = os.path.join(input_dir, case)
    combined_mask = None

    # Process each label in the case directory
    for label_name, label_value in label_dict.items():
        label_path = os.path.join(case_path, f"{label_name}.nrrd")

        if os.path.exists(label_path):
            # Load nrrd file
            data, header = nrrd.read(label_path)

            if combined_mask is None:
                combined_mask = np.zeros(data.shape, dtype=np.int16)

            # Apply the label value to the mask
            combined_mask[data > 0] = label_value

    # save the combined mask in Decathlon format (.nii.gz)
    if combined_mask is not None:
        output_path = os.path.join(output_dir, f"{case}.nii.gz")
        nii_img = nib.Nifti1Image(combined_mask, np.eye(4))
        nib.save(nii_img, output_path)
        print(f"Saved combined mask for {case} at {output_path}")
    else:
        print(f"No labels found for case {case}")`
  1. Restructure outputs in Decathlon format (or you can do this in step 2)

FYI: I didn't check if this code works or not it's just the example of the logic you are asking for

Mai-CS commented 2 weeks ago

Great! Solved the problem.