docker_task_2/nnsyn/nnunetv2/dataset_conversion/Dataset114_MNMs.py

import csv
import os
import random
from pathlib import Path

import nibabel as nib
from batchgenerators.utilities.file_and_folder_operations import load_json, save_json

from nnunetv2.dataset_conversion.Dataset027_ACDC import make_out_dirs
from nnunetv2.dataset_conversion.generate_dataset_json import generate_dataset_json
from nnunetv2.paths import nnUNet_preprocessed


def read_csv(csv_file: str):
    patient_info = {}

    with open(csv_file) as csvfile:
        reader = csv.reader(csvfile)
        headers = next(reader)
        patient_index = headers.index("External code")
        ed_index = headers.index("ED")
        es_index = headers.index("ES")
        vendor_index = headers.index("Vendor")

        for row in reader:
            patient_info[row[patient_index]] = {
                "ed": int(row[ed_index]),
                "es": int(row[es_index]),
                "vendor": row[vendor_index],
            }

    return patient_info


# ------------------------------------------------------------------------------
# Conversion to nnUNet format
# ------------------------------------------------------------------------------
def convert_mnms(src_data_folder: Path, csv_file_name: str, dataset_id: int):
    out_dir, out_train_dir, out_labels_dir, out_test_dir = make_out_dirs(dataset_id, task_name="MNMs")
    patients_train = [f for f in (src_data_folder / "Training" / "Labeled").iterdir() if f.is_dir()]
    patients_test = [f for f in (src_data_folder / "Testing").iterdir() if f.is_dir()]

    patient_info = read_csv(str(src_data_folder / csv_file_name))

    save_cardiac_phases(patients_train, patient_info, out_train_dir, out_labels_dir)
    save_cardiac_phases(patients_test, patient_info, out_test_dir)

    # There are non-orthonormal direction cosines in the test and validation data.
    # Not sure if the data should be fixed, or we should skip the problematic data.
    # patients_val = [f for f in (src_data_folder / "Validation").iterdir() if f.is_dir()]
    # save_cardiac_phases(patients_val, patient_info, out_train_dir, out_labels_dir)

    generate_dataset_json(
        str(out_dir),
        channel_names={
            0: "cineMRI",
        },
        labels={"background": 0, "LVBP": 1, "LVM": 2, "RV": 3},
        file_ending=".nii.gz",
        num_training_cases=len(patients_train) * 2,  # 2 since we have ED and ES for each patient
    )


def save_cardiac_phases(
    patients: list[Path], patient_info: dict[str, dict[str, int]], out_dir: Path, labels_dir: Path = None
):
    for patient in patients:
        print(f"Processing patient: {patient.name}")

        image = nib.load(patient / f"{patient.name}_sa.nii.gz")
        ed_frame = patient_info[patient.name]["ed"]
        es_frame = patient_info[patient.name]["es"]

        save_extracted_nifti_slice(image, ed_frame=ed_frame, es_frame=es_frame, out_dir=out_dir, patient=patient)

        if labels_dir:
            label = nib.load(patient / f"{patient.name}_sa_gt.nii.gz")
            save_extracted_nifti_slice(label, ed_frame=ed_frame, es_frame=es_frame, out_dir=labels_dir, patient=patient)


def save_extracted_nifti_slice(image, ed_frame: int, es_frame: int, out_dir: Path, patient: Path):
    # Save only extracted diastole and systole slices from the 4D H x W x D x time volume.
    image_ed = nib.Nifti1Image(image.dataobj[..., ed_frame], image.affine)
    image_es = nib.Nifti1Image(image.dataobj[..., es_frame], image.affine)

    # Labels do not have modality identifiers. Labels always end with 'gt'.
    suffix = ".nii.gz" if image.get_filename().endswith("_gt.nii.gz") else "_0000.nii.gz"

    nib.save(image_ed, str(out_dir / f"{patient.name}_frame{ed_frame:02d}{suffix}"))
    nib.save(image_es, str(out_dir / f"{patient.name}_frame{es_frame:02d}{suffix}"))


# ------------------------------------------------------------------------------
# Create custom splits
# ------------------------------------------------------------------------------
def create_custom_splits(src_data_folder: Path, csv_file: str, dataset_id: int, num_val_patients: int = 25):
    existing_splits = os.path.join(nnUNet_preprocessed, f"Dataset{dataset_id}_MNMs", "splits_final.json")
    splits = load_json(existing_splits)

    patients_train = [f.name for f in (src_data_folder / "Training" / "Labeled").iterdir() if f.is_dir()]
    # Filter out any patients not in the training set
    patient_info = {
        patient: data
        for patient, data in read_csv(str(src_data_folder / csv_file)).items()
        if patient in patients_train
    }

    # Get train and validation patients for both vendors
    patients_a = [patient for patient, patient_data in patient_info.items() if patient_data["vendor"] == "A"]
    patients_b = [patient for patient, patient_data in patient_info.items() if patient_data["vendor"] == "B"]
    train_a, val_a = get_vendor_split(patients_a, num_val_patients)
    train_b, val_b = get_vendor_split(patients_b, num_val_patients)

    # Build filenames from corresponding patient frames
    train_a = [f"{patient}_frame{patient_info[patient][frame]:02d}" for patient in train_a for frame in ["es", "ed"]]
    train_b = [f"{patient}_frame{patient_info[patient][frame]:02d}" for patient in train_b for frame in ["es", "ed"]]
    train_a_mix_1, train_a_mix_2 = train_a[: len(train_a) // 2], train_a[len(train_a) // 2 :]
    train_b_mix_1, train_b_mix_2 = train_b[: len(train_b) // 2], train_b[len(train_b) // 2 :]
    val_a = [f"{patient}_frame{patient_info[patient][frame]:02d}" for patient in val_a for frame in ["es", "ed"]]
    val_b = [f"{patient}_frame{patient_info[patient][frame]:02d}" for patient in val_b for frame in ["es", "ed"]]

    for train_set in [train_a, train_b, train_a_mix_1 + train_b_mix_1, train_a_mix_2 + train_b_mix_2]:
        # For each train set, we evaluate on A, B and (A + B) respectively
        # See table 3 from the original paper for more details.
        splits.append({"train": train_set, "val": val_a})
        splits.append({"train": train_set, "val": val_b})
        splits.append({"train": train_set, "val": val_a + val_b})

    save_json(splits, existing_splits)


def get_vendor_split(patients: list[str], num_val_patients: int):
    random.shuffle(patients)
    total_patients = len(patients)
    num_training_patients = total_patients - num_val_patients
    return patients[:num_training_patients], patients[num_training_patients:]


if __name__ == "__main__":
    import argparse

    class RawTextArgumentDefaultsHelpFormatter(argparse.ArgumentDefaultsHelpFormatter, argparse.RawTextHelpFormatter):
        pass

    parser = argparse.ArgumentParser(add_help=False, formatter_class=RawTextArgumentDefaultsHelpFormatter)
    parser.add_argument(
        "-h",
        "--help",
        action="help",
        default=argparse.SUPPRESS,
        help="MNMs conversion utility helper. This script can be used to convert MNMs data into the expected nnUNet "
        "format. It can also be used to create additional custom splits, for explicitly training on combinations "
        "of vendors A and B (see `--custom-splits`).\n"
        "If you wish to generate the custom splits, run the following pipeline:\n\n"
        "(1) Run `Dataset114_MNMs -i <raw_Data_dir>\n"
        "(2) Run `nnUNetv2_plan_and_preprocess -d 114 --verify_dataset_integrity`\n"
        "(3) Start training, but stop after initial splits are created: `nnUNetv2_train 114 2d 0`\n"
        "(4) Re-run `Dataset114_MNMs`, with `-s True`.\n"
        "(5) Re-run training.\n",
    )
    parser.add_argument(
        "-i",
        "--input_folder",
        type=str,
        default="./data/M&Ms/OpenDataset/",
        help="The downloaded MNMs dataset dir. Should contain a csv file, as well as Training, Validation and Testing "
        "folders.",
    )
    parser.add_argument(
        "-c",
        "--csv_file_name",
        type=str,
        default="211230_M&Ms_Dataset_information_diagnosis_opendataset.csv",
        help="The csv file containing the dataset information.",
    ),
    parser.add_argument("-d", "--dataset_id", type=int, default=114, help="nnUNet Dataset ID.")
    parser.add_argument(
        "-s",
        "--custom_splits",
        type=bool,
        default=False,
        help="Whether to append custom splits for training and testing on different vendors. If True, will create "
        "splits for training on patients from vendors A, B or a mix of A and B. Splits are tested on a hold-out "
        "validation sets of patients from A, B or A and B combined. See section 2.4 and table 3 from "
        "https://arxiv.org/abs/2011.07592 for more info.",
    )

    args = parser.parse_args()
    args.input_folder = Path(args.input_folder)

    if args.custom_splits:
        print("Appending custom splits...")
        create_custom_splits(args.input_folder, args.csv_file_name, args.dataset_id)
    else:
        print("Converting...")
        convert_mnms(args.input_folder, args.csv_file_name, args.dataset_id)

    print("Done!")