tutorials/BraTS/utils.py at f3a44648ea054224f18b8eaf4e0311f0a13e601b · BrainLesion/tutorials · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
from pathlib import Path
import matplotlib.pyplot as plt
import nibabel as nib

DATA_FOLDER = "data"


def visualize_segmentation_data(
    data_folder: str = DATA_FOLDER,
    subject_id: str = "BraTS-GLI-00001-000",
    slice_index: int = 75,
):
    """Visualize the MRI modalities for a given slice index

    Args:
        data_folder (str, optional): Path to the folder containing the t1, t1c, t2 & flair file. Defaults to DATA_FOLDER.
        slice_index (int, optional): Slice to be visualized (first index in data of shape (155, 240, 240)). Defaults to 75.
    """
    _, axes = plt.subplots(1, 4, figsize=(12, 10))

    subject_path = Path(data_folder) / subject_id
    modalities = ["t1n", "t1c", "t2f", "t2w"]
    for i, mod in enumerate(modalities):
        modality_file = subject_path / f"{subject_id}-{mod}.nii.gz"
        modality_np = nib.load(modality_file).get_fdata().transpose(2, 1, 0)
        axes[i].set_title(mod)
        axes[i].imshow(modality_np[slice_index, :, :], cmap="gray")
        axes[i].axis("off")


def visualize_inpainting_data(
    data_folder: str = DATA_FOLDER,
    subject_id: str = "BraTS-GLI-00001-000",
    slice_index: int = 75,
):
    """Visualize the MRI modalities for a given slice index

    Args:
        data_folder (str, optional): Path to the folder containing the t1n and mask files. Defaults to DATA_FOLDER.
        slice_index (int, optional): Slice to be visualized (first index in data of shape (155, 240, 240)). Defaults to 75.
    """
    _, axes = plt.subplots(1, 2, figsize=(6, 10))

    subject_path = Path(data_folder) / subject_id
    modalities = ["t1n-voided", "mask"]
    for i, mod in enumerate(modalities):
        modality_file = subject_path / f"{subject_id}-{mod}.nii.gz"
        modality_np = nib.load(modality_file).get_fdata().transpose(2, 1, 0)
        axes[i].set_title(mod)
        axes[i].imshow(modality_np[slice_index, :, :], cmap="gray")
        axes[i].axis("off")


def visualize_segmentation(modality_file: str, segmentation_file: str):
    """Visualize the MRI modality and the segmentation

    Args:
        modality_file (str): Path to the desired modality file
        segmentation_file (str): Path to the segmentation file
    """
    modality_np = nib.load(modality_file).get_fdata().transpose(2, 1, 0)
    seg_np = nib.load(segmentation_file).get_fdata().transpose(2, 1, 0)
    _, ax = plt.subplots(1, 2, figsize=(8, 4))

    slice_index = modality_np.shape[0] // 2  # You can choose any slice here
    ax[0].imshow(modality_np[slice_index, :, :], cmap="gray")
    ax[1].imshow(modality_np[slice_index, :, :], cmap="gray")
    ax[1].imshow(seg_np[slice_index, :, :], cmap="plasma", alpha=0.3)
    for ax in ax:
        ax.axis("off")
    plt.tight_layout()


def visualize_inpainting(t1n_voided: str, prediction: str):
    """Visualize the inpainting results

    Args:
        t1n_voided (str): Voided T1 modality file
        prediction (str): Inpainting prediction file
    """
    voided_np = nib.load(t1n_voided).get_fdata().transpose(2, 1, 0)
    inpainting_np = nib.load(prediction).get_fdata().transpose(2, 1, 0)
    _, ax = plt.subplots(1, 2, figsize=(8, 4))

    slice_index = voided_np.shape[0] // 2  # You can choose any slice here
    ax[0].imshow(voided_np[slice_index, :, :], cmap="gray")
    ax[1].imshow(inpainting_np[slice_index, :, :], cmap="gray")
    for ax in ax:
        ax.axis("off")
    plt.tight_layout()


def visualize_missing_mri_t2w(
    synthesized_t2w: str,
    data_folder: str = DATA_FOLDER,
    subject_id: str = "BraTS-GLI-00001-000",
    slice_index: int = 75,
):
    """Visualize the MRI modalities for a given slice index

    Args:
        data_folder (str, optional): Path to the folder containing the t1, t1c, t2 & flair file. Defaults to DATA_FOLDER.
        slice_index (int, optional): Slice to be visualized (first index in data of shape (155, 240, 240)). Defaults to 75.
    """
    _, axes = plt.subplots(1, 5, figsize=(12, 10))

    subject_path = Path(data_folder) / subject_id
    modalities = ["t1n", "t1c", "t2f", "t2w"]
    for i, mod in enumerate(modalities):
        modality_file = subject_path / f"{subject_id}-{mod}.nii.gz"
        modality_np = nib.load(modality_file).get_fdata().transpose(2, 1, 0)
        axes[i].set_title(mod)
        axes[i].imshow(modality_np[slice_index, :, :], cmap="gray")
        axes[i].axis("off")

    # show synthetic T2w
    synthetic_t2w_np = nib.load(synthesized_t2w).get_fdata().transpose(2, 1, 0)
    axes[4].set_title("Synthesized t2w")
    axes[4].imshow(synthetic_t2w_np[slice_index, :, :], cmap="gray")
    axes[4].axis("off")