Dataset119_ToothFairy2_All.py

from typing import Dict, Any
import os
from os.path import join
import json
import random
import multiprocessing

import SimpleITK as sitk
import numpy as np
from tqdm import tqdm


def mapping_DS119() -> Dict[int, int]:
    """Remove all NA Classes and make Class IDs continuous"""
    mapping = {}
    mapping.update({i: i for i in range(1, 19)})  # [1-10]->[1-10] | [11-18]->[11-18]
    mapping.update({i: i - 2 for i in range(21, 29)})  # [21-28]->[19-26]
    mapping.update({i: i - 4 for i in range(31, 39)})  # [31-38]->[27-34]
    mapping.update({i: i - 6 for i in range(41, 49)})  # [41-48]->[35-42]
    return mapping


def mapping_DS120() -> Dict[int, int]:
    """Remove Only Keep Teeth and Jaw Classes"""
    mapping = {}
    mapping.update({i: i for i in range(1, 3)})  # [0-2] -> [0-2]
    mapping.update({i: i - 8 for i in range(11, 19)})  # [11-18]->[3-10]
    mapping.update({i: i - 10 for i in range(21, 29)})  # [21-28]->[11-18]
    mapping.update({i: i - 12 for i in range(31, 39)})  # [31-38]->[19-26]
    mapping.update({i: i - 14 for i in range(41, 49)})  # [41-48]->[27-34]
    return mapping


def mapping_DS121() -> Dict[int, int]:
    """Remove Only Keep Teeth and Jaw Classes"""
    mapping = {}
    mapping.update({i: i - 10 for i in range(11, 19)})  # [11-18]->[3-8]
    mapping.update({i: i - 12 for i in range(21, 29)})  # [21-28]->[11-16]
    mapping.update({i: i - 14 for i in range(31, 39)})  # [31-38]->[19-24]
    mapping.update({i: i - 16 for i in range(41, 49)})  # [41-48]->[27-32]
    return mapping


def load_json(json_file: str) -> Any:
    with open(json_file, "r") as f:
        data = json.load(f)
    return data


def write_json(json_file: str, data: Any, indent: int = 4) -> None:
    with open(json_file, "w") as f:
        json.dump(data, f, indent=indent)


def image_to_nifi(input_path: str, output_path: str) -> None:
    image_sitk = sitk.ReadImage(input_path)
    sitk.WriteImage(image_sitk, output_path)


def label_mapping(input_path: str, output_path: str, mapping: Dict[int, int] = None) -> None:

    label_sitk = sitk.ReadImage(input_path)
    if mapping is not None:
        label_np = sitk.GetArrayFromImage(label_sitk)

        label_np_new = np.zeros_like(label_np, dtype=np.uint8)
        for org_id, new_id in mapping.items():
            label_np_new[label_np == org_id] = new_id

        label_sitk_new = sitk.GetImageFromArray(label_np_new)
        label_sitk_new.CopyInformation(label_sitk)
        sitk.WriteImage(label_sitk_new, output_path)
    else:
        sitk.WriteImage(label_sitk, output_path)


def process_images(files: str, img_dir_in: str, img_dir_out: str, n_processes: int = 12):

    os.makedirs(img_dir_out, exist_ok=True)

    iterable = [
        {
            "input_path": join(img_dir_in, file),
            "output_path": join(img_dir_out, file.replace(".mha", ".nii.gz")),
        }
        for file in files
    ]
    with multiprocessing.Pool(processes=n_processes) as pool:
        jobs = [pool.apply_async(image_to_nifi, kwds={**args}) for args in iterable]
        _ = [job.get() for job in tqdm(jobs, desc="Process Images")]


def process_labels(
    files: str, lbl_dir_in: str, lbl_dir_out: str, mapping: Dict[int, int], n_processes: int = 12
) -> None:

    os.makedirs(lbl_dir_out, exist_ok=True)

    iterable = [
        {
            "input_path": join(lbl_dir_in, file),
            "output_path": join(lbl_dir_out, file.replace(".mha", ".nii.gz")),
            "mapping": mapping,
        }
        for file in files
    ]
    with multiprocessing.Pool(processes=n_processes) as pool:
        jobs = [pool.apply_async(label_mapping, kwds={**args}) for args in iterable]
        _ = [job.get() for job in tqdm(jobs, desc="Process Labels...")]


def process_ds(
    root: str, input_ds: str, output_ds: str, mapping: dict, image_link: str = None
) -> None:
    os.makedirs(join(root, output_ds), exist_ok=True)
    os.makedirs(join(root, output_ds, "labelsTr"), exist_ok=True)
    # --- Handle Labels --- #
    lbl_files = os.listdir(join(root, input_ds, "labelsTr"))
    lbl_dir_in = join(root, input_ds, "labelsTr")
    lbl_dir_out = join(root, output_ds, "labelsTr")

    process_labels(lbl_files, lbl_dir_in, lbl_dir_out, mapping, n_processes=12)

    # --- Handle Images --- #
    img_files = os.listdir(join(root, input_ds, "imagesTr"))
    dataset = {}
    if image_link is None:
        img_dir_in = join(root, input_ds, "imagesTr")
        img_dir_out = join(root, output_ds, "imagesTr")

        process_images(img_files, img_dir_in, img_dir_out, n_processes=12)
    else:
        base_name = [file.replace("_0000.mha", "") for file in img_files]
        for name in base_name:
            dataset[name] = {
                "images": [join("..", image_link, "imagesTr", name + "_0000.nii.gz")],
                "label": join("labelsTr", name + ".nii.gz"),
            }

    # --- Generate dataset.json --- #
    dataset_json = load_json(join(root, input_ds, "dataset.json"))
    dataset_json["file_ending"] = ".nii.gz"
    dataset_json["name"] = output_ds
    dataset_json["numTraining"] = len(lbl_files)
    if dataset != {}:
        dataset_json["dataset"] = dataset

    label_dict = dataset_json["labels"]
    label_dict_new = {"background": 0}
    for k, v in label_dict.items():
        if v in mapping.keys():
            label_dict_new[k] = mapping[v]
    dataset_json["labels"] = label_dict_new
    write_json(join(root, output_ds, "dataset.json"), dataset_json)

    # --- Generate splits_final.json --- #
    img_names = [file.replace("_0000.mha", "") for file in img_files]

    random_seed = 42
    random.seed(random_seed)
    random.shuffle(img_names)

    split_index = int(len(img_names) * 0.7)  # 70:30 split
    train_files = img_names[:split_index]
    val_files = img_names[split_index:]
    train_files.sort()
    val_files.sort()

    split = [{"train": train_files, "val": val_files}]
    write_json(join(root, output_ds, "splits_final.json"), split)


if __name__ == "__main__":
    # Different nnUNet Datasets
    # Dataset 112: Raw
    # Dataset 119: Replace NaN classes
    # Dataset 120: Only Teeth + Jaw Classes
    # Dataset 121: Only Teeth Classes

    root = "/media/l727r/data/Teeth_Data/ToothFairy2_Dataset"

    process_ds(root, "Dataset112_ToothFairy2", "Dataset119_ToothFairy2_All", mapping_DS119(), None)
    # process_ds(
    #     root,
    #     "Dataset112_ToothFairy2",
    #     "Dataset120_ToothFairy2_JawTeeth",
    #     mapping_DS120(),
    #     "Dataset119_ToothFairy2_All",
    # )
    # process_ds(
    #     root,
    #     "Dataset112_ToothFairy2",
    #     "Dataset121_ToothFairy2_Teeth",
    #     mapping_DS121(),
    #     "Dataset119_ToothFairy2_All",
    # )