inference.postproc.postprocess

Post-processing

Provides functionalities related to computing connected components and postprocessing segmentation predictions by removing small connected components based on thresholds.

  1"""
  2### Post-processing 
  3
  4Provides functionalities related to computing connected components and postprocessing segmentation predictions by removing small connected components based on thresholds.
  5"""
  6
  7
  8from pathlib import Path
  9import pandas as pd
 10import pickle
 11import argparse
 12import json
 13import numpy as np
 14import nibabel as nib
 15import cc3d
 16import SimpleITK as sitk
 17import glob
 18import os
 19from tqdm import tqdm
 20import shutil
 21import warnings
 22from typing import Any, List, Tuple, Dict
 23
 24
 25def maybe_make_dir(path: str) -> None:
 26    """
 27    Creates a directory at the specified path if it does not exist.
 28
 29    Args:
 30        path (str): A string representing the directory path.
 31    """
 32    if not os.path.isdir(path):
 33        os.makedirs(path)
 34
 35
 36def read_image(path: Path) -> Tuple[sitk.Image, np.ndarray]:
 37    """
 38    Reads an image file and returns the SimpleITK image object and its NumPy array.
 39
 40    Args:
 41        path (Path): Path to the image file.
 42
 43    Returns:
 44        Tuple[sitk.Image, np.ndarray]: The SimpleITK image object and its corresponding NumPy array.
 45    """
 46    img_sitk = sitk.ReadImage(path)
 47    img = sitk.GetArrayFromImage(img_sitk)
 48    return img_sitk, img
 49
 50
 51def convert_labels_back_to_BraTS(seg: np.ndarray) -> np.ndarray:
 52    """
 53    Converts segmentation labels back to BraTS format.
 54
 55    Args:
 56        seg (np.ndarray): The segmentation array.
 57
 58    Returns:
 59        np.ndarray: The converted segmentation array.
 60    """
 61    new_seg = np.zeros_like(seg)
 62    new_seg[seg == 1] = 2
 63    new_seg[seg == 3] = 3
 64    new_seg[seg == 2] = 1
 65    return new_seg
 66
 67
 68def get_ratio_ncr_wt(seg: np.ndarray) -> float:
 69    """
 70    Calculates the ratio of necrotic and non-enhancing tumor core (NCR) voxels to whole tumor (WT) voxels.
 71
 72    Args:
 73        seg (np.ndarray): The segmentation array.
 74
 75    Returns:
 76        float: The NCR to WT voxel ratio.
 77    """
 78    ncr_voxels = np.sum(seg == 1)
 79    wt_voxels = np.sum(seg != 0)
 80    if wt_voxels == 0:
 81        return 1.0
 82    return ncr_voxels / wt_voxels
 83
 84
 85def get_ratio_ed_wt(seg: np.ndarray) -> float:
 86    """
 87    Calculates the ratio of peritumoral edema (ED) voxels to whole tumor (WT) voxels.
 88
 89    Args:
 90        seg (np.ndarray): The segmentation array.
 91
 92    Returns:
 93        float: The ED to WT voxel ratio.
 94    """
 95    ed_voxels = np.sum(seg == 2)
 96    wt_voxels = np.sum(seg != 0)
 97    if wt_voxels == 0:
 98        return 1.0
 99    return ed_voxels / wt_voxels
100
101
102def get_ratio_et_wt(seg: np.ndarray) -> float:
103    """
104    Calculates the ratio of enhancing tumor (ET) voxels to whole tumor (WT) voxels.
105
106    Args:
107        seg (np.ndarray): The segmentation array.
108
109    Returns:
110        float: The ET to WT voxel ratio.
111    """
112    et_voxels = np.sum(seg == 3)
113    wt_voxels = np.sum(seg != 0)
114    if wt_voxels == 0:
115        return 1.0
116    return et_voxels / wt_voxels
117
118
119def get_ratio_tc_wt(seg: np.ndarray) -> float:
120    """
121    Calculates the ratio of tumor core (TC) voxels to whole tumor (WT) voxels.
122
123    Args:
124        seg (np.ndarray): The segmentation array.
125
126    Returns:
127        float: The TC to WT voxel ratio.
128    """
129    tc_voxels = np.sum((seg == 1) & (seg == 3))
130    wt_voxels = np.sum(seg != 0)
131    if wt_voxels == 0:
132        return 1.0
133    return tc_voxels / wt_voxels
134
135
136def convert_et_to_ncr(seg: np.ndarray) -> np.ndarray:
137    """
138    Converts enhancing tumor (ET) labels to necrotic and non-enhancing tumor core (NCR).
139
140    Args:
141        seg (np.ndarray): The segmentation array.
142
143    Returns:
144        np.ndarray: The modified segmentation array.
145    """
146    seg[seg == 3] = 1
147    return seg
148
149
150def convert_ed_to_ncr(seg: np.ndarray) -> np.ndarray:
151    """
152    Converts peritumoral edema (ED) labels to necrotic and non-enhancing tumor core (NCR).
153
154    Args:
155        seg (np.ndarray): The segmentation array.
156
157    Returns:
158        np.ndarray: The modified segmentation array.
159    """
160    seg[seg == 2] = 1
161    return seg
162
163
164def get_greatest_label(seg: np.ndarray) -> Tuple[str, float]:
165    """
166    Determines the label with the highest ratio to whole tumor (WT) voxels.
167
168    Args:
169        seg (np.ndarray): The segmentation array.
170
171    Returns:
172        Tuple[str, float]: The label with the highest ratio and its corresponding ratio value.
173    """
174    ratios = {
175        "ncr": get_ratio_ncr_wt(seg),
176        "ed": get_ratio_ed_wt(seg),
177        "et": get_ratio_et_wt(seg),
178        # "tc": get_ratio_tc_wt(seg),
179    }
180    greatest_label = max(ratios, key=ratios.get)
181    return greatest_label, ratios[greatest_label]
182
183
184def redefine_et_ed_labels(
185    seg_file: Path,
186    out_file: Path,
187    label: str = "et",
188    ratio: float = 0.0
189) -> np.ndarray:
190    """
191    Redefines ET or ED labels to NCR in the segmentation based on a specified ratio.
192
193    Args:
194        seg_file (Path): Path to the input segmentation file.
195        out_file (Path): Path to save the postprocessed segmentation file.
196        label (str, optional): Label to optimize ("et" or "ed"). Defaults to "et".
197        ratio (float, optional): Threshold ratio for label optimization. Defaults to 0.0.
198
199    Returns:
200        np.ndarray: The modified segmentation array.
201    """
202    seg_obj = nib.load(seg_file)
203    seg = seg_obj.get_fdata()
204    if label == "et":
205        ratio_et_wt = get_ratio_et_wt(seg)
206        if ratio_et_wt < ratio:
207            seg = convert_et_to_ncr(seg)
208    elif label == "ed":
209        ratio_ed_wt = get_ratio_ed_wt(seg)
210        if ratio_ed_wt < ratio:
211            seg = convert_ed_to_ncr(seg)
212    new_obj = nib.Nifti1Image(seg.astype(np.int8), seg_obj.affine)
213    nib.save(new_obj, out_file)
214    return seg
215
216
217def postprocess_image(
218    seg: np.ndarray,
219    label: str,
220    ratio: float = 0.04
221) -> np.ndarray:
222    """
223    Postprocesses the segmentation image by redefining ET or ED labels based on a ratio.
224
225    Args:
226        seg (np.ndarray): The segmentation array.
227        label (str): Label to optimize ("et" or "ed").
228        ratio (float, optional): Threshold ratio for label optimization. Defaults to 0.04.
229
230    Returns:
231        np.ndarray: The postprocessed segmentation array.
232    """
233    if label == "et":
234        ratio_et_wt = get_ratio_et_wt(seg)
235        if ratio_et_wt < ratio:
236            seg = convert_et_to_ncr(seg)
237    elif label == "ed":
238        ratio_ed_wt = get_ratio_ed_wt(seg)
239        if ratio_ed_wt < ratio:
240            seg = convert_ed_to_ncr(seg)
241    return seg
242
243
244def save_image(
245    img: np.ndarray,
246    img_sitk: sitk.Image,
247    out_path: Path
248) -> None:
249    """
250    Saves the NumPy array as a NIfTI image with the original image's metadata.
251
252    Args:
253        img (np.ndarray): The image array to save.
254        img_sitk (sitk.Image): The original SimpleITK image object.
255        out_path (Path): Path to save the new image.
256    """
257    new_img_sitk = sitk.GetImageFromArray(img)
258    new_img_sitk.CopyInformation(img_sitk)
259    sitk.WriteImage(new_img_sitk, out_path)
260
261
262def postprocess_batch(
263    input_folder: Path,
264    output_folder: Path,
265    label_to_optimize: str,
266    ratio: float = 0.04,
267    convert_to_brats_labels: bool = False
268) -> None:
269    """
270    Postprocesses a batch of segmentation files by optimizing specified labels.
271
272    Args:
273        input_folder (Path): Path to the input directory containing segmentation files.
274        output_folder (Path): Path to the output directory to save postprocessed files.
275        label_to_optimize (str): Label to optimize ("et" or "ed").
276        ratio (float, optional): Threshold ratio for label optimization. Defaults to 0.04.
277        convert_to_brats_labels (bool, optional): Whether to convert labels back to BraTS format. Defaults to False.
278    """
279    seg_list = sorted(glob.glob(os.path.join(input_folder, "*.nii.gz")))
280    for seg_path in tqdm(seg_list):
281        seg_sitk, seg = read_image(Path(seg_path))
282        if convert_to_brats_labels:
283            seg = convert_labels_back_to_BraTS(seg)
284        seg_pp = postprocess_image(seg, label_to_optimize, ratio)
285        out_path = output_folder / Path(seg_path).name
286        save_image(seg_pp, seg_sitk, out_path)
287
288
289def get_connected_labels(seg_file: Path) -> Tuple[np.ndarray, np.ndarray, np.ndarray, int, int, int]:
290    """
291    Identifies connected components for NCR, ED, and ET labels in the segmentation.
292
293    Args:
294        seg_file (Path): Path to the segmentation file.
295
296    Returns:
297        Tuple[np.ndarray, np.ndarray, np.ndarray, int, int, int]: 
298            - Connected labels for NCR, ED, ET
299            - Number of connected components for NCR, ED, ET
300    """
301    seg_obj = nib.load(seg_file)
302    seg = seg_obj.get_fdata()
303    seg_ncr = np.where(seg == 1, 1, 0)
304    seg_ed = np.where(seg == 2, 2, 0)
305    seg_et = np.where(seg == 3, 3, 0)
306    labels_ncr, n_ncr = cc3d.connected_components(seg_ncr, connectivity=26, return_N=True)
307    labels_ed, n_ed = cc3d.connected_components(seg_ed, connectivity=26, return_N=True)
308    labels_et, n_et = cc3d.connected_components(seg_et, connectivity=26, return_N=True)
309    return labels_ncr, labels_ed, labels_et, n_ncr, n_ed, n_et
310
311
312def remove_disconnected(
313    seg_file: Path,
314    out_file: Path,
315    t_ncr: int = 50,
316    t_ed: int = 50,
317    t_et: int = 50
318) -> Tuple[int, int, int, int, int, int]:
319    """
320    Removes disconnected small components from the segmentation based on thresholds.
321
322    Args:
323        seg_file (Path): Path to the input segmentation file.
324        out_file (Path): Path to save the cleaned segmentation file.
325        t_ncr (int, optional): Threshold for NCR voxel count. Defaults to 50.
326        t_ed (int, optional): Threshold for ED voxel count. Defaults to 50.
327        t_et (int, optional): Threshold for ET voxel count. Defaults to 50.
328
329    Returns:
330        Tuple[int, int, int, int, int, int]: 
331            Number of removed NCR, total NCR, removed ED, total ED, removed ET, total ET.
332    """
333    seg_obj = nib.load(seg_file)
334    labels_ncr, labels_ed, labels_et, n_ncr, n_ed, n_et = get_connected_labels(seg_file)
335    
336    # Process NCR
337    vols = []
338    for i in range(n_ncr):
339        tmp = np.where(labels_ncr == i + 1, 1, 0)
340        vol = np.count_nonzero(tmp)
341        if vol < t_ncr:
342            labels_ncr = np.where(labels_ncr == i + 1, 0, labels_ncr)
343            vols.append(vol)
344    removed_ncr = len(vols)
345    
346    # Process ED
347    vols = []
348    for i in range(n_ed):
349        tmp = np.where(labels_ed == i + 1, 1, 0)
350        vol = np.count_nonzero(tmp)
351        if vol < t_ed:
352            labels_ed = np.where(labels_ed == i + 1, 0, labels_ed)
353            vols.append(vol)
354    removed_ed = len(vols)
355    
356    # Process ET
357    vols = []
358    for i in range(n_et):
359        tmp = np.where(labels_et == i + 1, 1, 0)
360        vol = np.count_nonzero(tmp)
361        if vol < t_et:
362            labels_et = np.where(labels_et == i + 1, 0, labels_et)
363            vols.append(vol)
364    removed_et = len(vols)
365
366    # Combine the cleaned labels
367    new_ncr = np.where(labels_ncr != 0, 1, 0)
368    new_ed = np.where(labels_ed != 0, 2, 0)
369    new_et = np.where(labels_et != 0, 3, 0)
370    new_seg = new_ncr + new_ed + new_et
371    new_obj = nib.Nifti1Image(new_seg.astype(np.int8), seg_obj.affine)
372    nib.save(new_obj, out_file)
373    return removed_ncr, n_ncr, removed_ed, n_ed, removed_et, n_et
374
375
376def remove_disconnected_from_dir(
377    input_dir: Path,
378    output_dir: Path,
379    t_ncr: int = 50,
380    t_ed: int = 50,
381    t_et: int = 50
382) -> Path:
383    """
384    Removes disconnected small components from all segmentation files in a directory.
385
386    Args:
387        input_dir (Path): Path to the input directory containing segmentation files.
388        output_dir (Path): Path to the output directory to save cleaned segmentation files.
389        t_ncr (int, optional): Threshold for NCR voxel count. Defaults to 50.
390        t_ed (int, optional): Threshold for ED voxel count. Defaults to 50.
391        t_et (int, optional): Threshold for ET voxel count. Defaults to 50.
392
393    Returns:
394        Path: Path to the output directory.
395    """
396    if not output_dir.exists():
397        output_dir.mkdir(parents=True)
398
399    for seg1 in input_dir.iterdir():
400        if seg1.name.endswith('.nii.gz'):
401            casename = seg1.name
402            seg2 = output_dir / casename
403            removed_ncr, n_ncr, removed_ed, n_ed, removed_et, n_et = remove_disconnected(
404                seg1, seg2, t_ncr=t_ncr, t_ed=t_ed, t_et=t_et
405            )
406            print(
407                f"{casename} removed regions NCR {removed_ncr:03d}/{n_ncr:03d} "
408                f"ED {removed_ed:03d}/{n_ed:03d} ET {removed_et:03d}/{n_et:03d}"
409            )
410        else:
411            print("Wrong input file!")
412
413    return output_dir
def maybe_make_dir(path: str) -> None: 
26def maybe_make_dir(path: str) -> None:
27    """
28    Creates a directory at the specified path if it does not exist.
29
30    Args:
31        path (str): A string representing the directory path.
32    """
33    if not os.path.isdir(path):
34        os.makedirs(path)

Creates a directory at the specified path if it does not exist.

Args: path (str): A string representing the directory path.

def read_image(path: pathlib.Path) -> Tuple[SimpleITK.SimpleITK.Image, numpy.ndarray]: 
37def read_image(path: Path) -> Tuple[sitk.Image, np.ndarray]:
38    """
39    Reads an image file and returns the SimpleITK image object and its NumPy array.
40
41    Args:
42        path (Path): Path to the image file.
43
44    Returns:
45        Tuple[sitk.Image, np.ndarray]: The SimpleITK image object and its corresponding NumPy array.
46    """
47    img_sitk = sitk.ReadImage(path)
48    img = sitk.GetArrayFromImage(img_sitk)
49    return img_sitk, img

Reads an image file and returns the SimpleITK image object and its NumPy array.

Args: path (Path): Path to the image file.

Returns: Tuple[sitk.Image, np.ndarray]: The SimpleITK image object and its corresponding NumPy array.

def convert_labels_back_to_BraTS(seg: numpy.ndarray) -> numpy.ndarray: 
52def convert_labels_back_to_BraTS(seg: np.ndarray) -> np.ndarray:
53    """
54    Converts segmentation labels back to BraTS format.
55
56    Args:
57        seg (np.ndarray): The segmentation array.
58
59    Returns:
60        np.ndarray: The converted segmentation array.
61    """
62    new_seg = np.zeros_like(seg)
63    new_seg[seg == 1] = 2
64    new_seg[seg == 3] = 3
65    new_seg[seg == 2] = 1
66    return new_seg

Converts segmentation labels back to BraTS format.

Args: seg (np.ndarray): The segmentation array.

Returns: np.ndarray: The converted segmentation array.

def get_ratio_ncr_wt(seg: numpy.ndarray) -> float: 
69def get_ratio_ncr_wt(seg: np.ndarray) -> float:
70    """
71    Calculates the ratio of necrotic and non-enhancing tumor core (NCR) voxels to whole tumor (WT) voxels.
72
73    Args:
74        seg (np.ndarray): The segmentation array.
75
76    Returns:
77        float: The NCR to WT voxel ratio.
78    """
79    ncr_voxels = np.sum(seg == 1)
80    wt_voxels = np.sum(seg != 0)
81    if wt_voxels == 0:
82        return 1.0
83    return ncr_voxels / wt_voxels

Calculates the ratio of necrotic and non-enhancing tumor core (NCR) voxels to whole tumor (WT) voxels.

Args: seg (np.ndarray): The segmentation array.

Returns: float: The NCR to WT voxel ratio.

def get_ratio_ed_wt(seg: numpy.ndarray) -> float: 
 86def get_ratio_ed_wt(seg: np.ndarray) -> float:
 87    """
 88    Calculates the ratio of peritumoral edema (ED) voxels to whole tumor (WT) voxels.
 89
 90    Args:
 91        seg (np.ndarray): The segmentation array.
 92
 93    Returns:
 94        float: The ED to WT voxel ratio.
 95    """
 96    ed_voxels = np.sum(seg == 2)
 97    wt_voxels = np.sum(seg != 0)
 98    if wt_voxels == 0:
 99        return 1.0
100    return ed_voxels / wt_voxels

Calculates the ratio of peritumoral edema (ED) voxels to whole tumor (WT) voxels.

Args: seg (np.ndarray): The segmentation array.

Returns: float: The ED to WT voxel ratio.

def get_ratio_et_wt(seg: numpy.ndarray) -> float: 
103def get_ratio_et_wt(seg: np.ndarray) -> float:
104    """
105    Calculates the ratio of enhancing tumor (ET) voxels to whole tumor (WT) voxels.
106
107    Args:
108        seg (np.ndarray): The segmentation array.
109
110    Returns:
111        float: The ET to WT voxel ratio.
112    """
113    et_voxels = np.sum(seg == 3)
114    wt_voxels = np.sum(seg != 0)
115    if wt_voxels == 0:
116        return 1.0
117    return et_voxels / wt_voxels

Calculates the ratio of enhancing tumor (ET) voxels to whole tumor (WT) voxels.

Args: seg (np.ndarray): The segmentation array.

Returns: float: The ET to WT voxel ratio.

def get_ratio_tc_wt(seg: numpy.ndarray) -> float: 
120def get_ratio_tc_wt(seg: np.ndarray) -> float:
121    """
122    Calculates the ratio of tumor core (TC) voxels to whole tumor (WT) voxels.
123
124    Args:
125        seg (np.ndarray): The segmentation array.
126
127    Returns:
128        float: The TC to WT voxel ratio.
129    """
130    tc_voxels = np.sum((seg == 1) & (seg == 3))
131    wt_voxels = np.sum(seg != 0)
132    if wt_voxels == 0:
133        return 1.0
134    return tc_voxels / wt_voxels

Calculates the ratio of tumor core (TC) voxels to whole tumor (WT) voxels.

Args: seg (np.ndarray): The segmentation array.

Returns: float: The TC to WT voxel ratio.

def convert_et_to_ncr(seg: numpy.ndarray) -> numpy.ndarray: 
137def convert_et_to_ncr(seg: np.ndarray) -> np.ndarray:
138    """
139    Converts enhancing tumor (ET) labels to necrotic and non-enhancing tumor core (NCR).
140
141    Args:
142        seg (np.ndarray): The segmentation array.
143
144    Returns:
145        np.ndarray: The modified segmentation array.
146    """
147    seg[seg == 3] = 1
148    return seg

Converts enhancing tumor (ET) labels to necrotic and non-enhancing tumor core (NCR).

Args: seg (np.ndarray): The segmentation array.

Returns: np.ndarray: The modified segmentation array.

def convert_ed_to_ncr(seg: numpy.ndarray) -> numpy.ndarray: 
151def convert_ed_to_ncr(seg: np.ndarray) -> np.ndarray:
152    """
153    Converts peritumoral edema (ED) labels to necrotic and non-enhancing tumor core (NCR).
154
155    Args:
156        seg (np.ndarray): The segmentation array.
157
158    Returns:
159        np.ndarray: The modified segmentation array.
160    """
161    seg[seg == 2] = 1
162    return seg

Converts peritumoral edema (ED) labels to necrotic and non-enhancing tumor core (NCR).

Args: seg (np.ndarray): The segmentation array.

Returns: np.ndarray: The modified segmentation array.

def get_greatest_label(seg: numpy.ndarray) -> Tuple[str, float]: 
165def get_greatest_label(seg: np.ndarray) -> Tuple[str, float]:
166    """
167    Determines the label with the highest ratio to whole tumor (WT) voxels.
168
169    Args:
170        seg (np.ndarray): The segmentation array.
171
172    Returns:
173        Tuple[str, float]: The label with the highest ratio and its corresponding ratio value.
174    """
175    ratios = {
176        "ncr": get_ratio_ncr_wt(seg),
177        "ed": get_ratio_ed_wt(seg),
178        "et": get_ratio_et_wt(seg),
179        # "tc": get_ratio_tc_wt(seg),
180    }
181    greatest_label = max(ratios, key=ratios.get)
182    return greatest_label, ratios[greatest_label]

Determines the label with the highest ratio to whole tumor (WT) voxels.

Args: seg (np.ndarray): The segmentation array.

Returns: Tuple[str, float]: The label with the highest ratio and its corresponding ratio value.

def redefine_et_ed_labels( seg_file: pathlib.Path, out_file: pathlib.Path, label: str = 'et', ratio: float = 0.0) -> numpy.ndarray: 
185def redefine_et_ed_labels(
186    seg_file: Path,
187    out_file: Path,
188    label: str = "et",
189    ratio: float = 0.0
190) -> np.ndarray:
191    """
192    Redefines ET or ED labels to NCR in the segmentation based on a specified ratio.
193
194    Args:
195        seg_file (Path): Path to the input segmentation file.
196        out_file (Path): Path to save the postprocessed segmentation file.
197        label (str, optional): Label to optimize ("et" or "ed"). Defaults to "et".
198        ratio (float, optional): Threshold ratio for label optimization. Defaults to 0.0.
199
200    Returns:
201        np.ndarray: The modified segmentation array.
202    """
203    seg_obj = nib.load(seg_file)
204    seg = seg_obj.get_fdata()
205    if label == "et":
206        ratio_et_wt = get_ratio_et_wt(seg)
207        if ratio_et_wt < ratio:
208            seg = convert_et_to_ncr(seg)
209    elif label == "ed":
210        ratio_ed_wt = get_ratio_ed_wt(seg)
211        if ratio_ed_wt < ratio:
212            seg = convert_ed_to_ncr(seg)
213    new_obj = nib.Nifti1Image(seg.astype(np.int8), seg_obj.affine)
214    nib.save(new_obj, out_file)
215    return seg

Redefines ET or ED labels to NCR in the segmentation based on a specified ratio.

Args: seg_file (Path): Path to the input segmentation file. out_file (Path): Path to save the postprocessed segmentation file. label (str, optional): Label to optimize ("et" or "ed"). Defaults to "et". ratio (float, optional): Threshold ratio for label optimization. Defaults to 0.0.

Returns: np.ndarray: The modified segmentation array.

def postprocess_image(seg: numpy.ndarray, label: str, ratio: float = 0.04) -> numpy.ndarray: 
218def postprocess_image(
219    seg: np.ndarray,
220    label: str,
221    ratio: float = 0.04
222) -> np.ndarray:
223    """
224    Postprocesses the segmentation image by redefining ET or ED labels based on a ratio.
225
226    Args:
227        seg (np.ndarray): The segmentation array.
228        label (str): Label to optimize ("et" or "ed").
229        ratio (float, optional): Threshold ratio for label optimization. Defaults to 0.04.
230
231    Returns:
232        np.ndarray: The postprocessed segmentation array.
233    """
234    if label == "et":
235        ratio_et_wt = get_ratio_et_wt(seg)
236        if ratio_et_wt < ratio:
237            seg = convert_et_to_ncr(seg)
238    elif label == "ed":
239        ratio_ed_wt = get_ratio_ed_wt(seg)
240        if ratio_ed_wt < ratio:
241            seg = convert_ed_to_ncr(seg)
242    return seg

Postprocesses the segmentation image by redefining ET or ED labels based on a ratio.

Args: seg (np.ndarray): The segmentation array. label (str): Label to optimize ("et" or "ed"). ratio (float, optional): Threshold ratio for label optimization. Defaults to 0.04.

Returns: np.ndarray: The postprocessed segmentation array.

def save_image( img: numpy.ndarray, img_sitk: SimpleITK.SimpleITK.Image, out_path: pathlib.Path) -> None: 
245def save_image(
246    img: np.ndarray,
247    img_sitk: sitk.Image,
248    out_path: Path
249) -> None:
250    """
251    Saves the NumPy array as a NIfTI image with the original image's metadata.
252
253    Args:
254        img (np.ndarray): The image array to save.
255        img_sitk (sitk.Image): The original SimpleITK image object.
256        out_path (Path): Path to save the new image.
257    """
258    new_img_sitk = sitk.GetImageFromArray(img)
259    new_img_sitk.CopyInformation(img_sitk)
260    sitk.WriteImage(new_img_sitk, out_path)

Saves the NumPy array as a NIfTI image with the original image's metadata.

Args: img (np.ndarray): The image array to save. img_sitk (sitk.Image): The original SimpleITK image object. out_path (Path): Path to save the new image.

def postprocess_batch( input_folder: pathlib.Path, output_folder: pathlib.Path, label_to_optimize: str, ratio: float = 0.04, convert_to_brats_labels: bool = False) -> None: 
263def postprocess_batch(
264    input_folder: Path,
265    output_folder: Path,
266    label_to_optimize: str,
267    ratio: float = 0.04,
268    convert_to_brats_labels: bool = False
269) -> None:
270    """
271    Postprocesses a batch of segmentation files by optimizing specified labels.
272
273    Args:
274        input_folder (Path): Path to the input directory containing segmentation files.
275        output_folder (Path): Path to the output directory to save postprocessed files.
276        label_to_optimize (str): Label to optimize ("et" or "ed").
277        ratio (float, optional): Threshold ratio for label optimization. Defaults to 0.04.
278        convert_to_brats_labels (bool, optional): Whether to convert labels back to BraTS format. Defaults to False.
279    """
280    seg_list = sorted(glob.glob(os.path.join(input_folder, "*.nii.gz")))
281    for seg_path in tqdm(seg_list):
282        seg_sitk, seg = read_image(Path(seg_path))
283        if convert_to_brats_labels:
284            seg = convert_labels_back_to_BraTS(seg)
285        seg_pp = postprocess_image(seg, label_to_optimize, ratio)
286        out_path = output_folder / Path(seg_path).name
287        save_image(seg_pp, seg_sitk, out_path)

Postprocesses a batch of segmentation files by optimizing specified labels.

Args: input_folder (Path): Path to the input directory containing segmentation files. output_folder (Path): Path to the output directory to save postprocessed files. label_to_optimize (str): Label to optimize ("et" or "ed"). ratio (float, optional): Threshold ratio for label optimization. Defaults to 0.04. convert_to_brats_labels (bool, optional): Whether to convert labels back to BraTS format. Defaults to False.

def get_connected_labels( seg_file: pathlib.Path) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray, int, int, int]: 
290def get_connected_labels(seg_file: Path) -> Tuple[np.ndarray, np.ndarray, np.ndarray, int, int, int]:
291    """
292    Identifies connected components for NCR, ED, and ET labels in the segmentation.
293
294    Args:
295        seg_file (Path): Path to the segmentation file.
296
297    Returns:
298        Tuple[np.ndarray, np.ndarray, np.ndarray, int, int, int]: 
299            - Connected labels for NCR, ED, ET
300            - Number of connected components for NCR, ED, ET
301    """
302    seg_obj = nib.load(seg_file)
303    seg = seg_obj.get_fdata()
304    seg_ncr = np.where(seg == 1, 1, 0)
305    seg_ed = np.where(seg == 2, 2, 0)
306    seg_et = np.where(seg == 3, 3, 0)
307    labels_ncr, n_ncr = cc3d.connected_components(seg_ncr, connectivity=26, return_N=True)
308    labels_ed, n_ed = cc3d.connected_components(seg_ed, connectivity=26, return_N=True)
309    labels_et, n_et = cc3d.connected_components(seg_et, connectivity=26, return_N=True)
310    return labels_ncr, labels_ed, labels_et, n_ncr, n_ed, n_et

Identifies connected components for NCR, ED, and ET labels in the segmentation.

Args: seg_file (Path): Path to the segmentation file.

Returns: Tuple[np.ndarray, np.ndarray, np.ndarray, int, int, int]: - Connected labels for NCR, ED, ET - Number of connected components for NCR, ED, ET

def remove_disconnected( seg_file: pathlib.Path, out_file: pathlib.Path, t_ncr: int = 50, t_ed: int = 50, t_et: int = 50) -> Tuple[int, int, int, int, int, int]: 
313def remove_disconnected(
314    seg_file: Path,
315    out_file: Path,
316    t_ncr: int = 50,
317    t_ed: int = 50,
318    t_et: int = 50
319) -> Tuple[int, int, int, int, int, int]:
320    """
321    Removes disconnected small components from the segmentation based on thresholds.
322
323    Args:
324        seg_file (Path): Path to the input segmentation file.
325        out_file (Path): Path to save the cleaned segmentation file.
326        t_ncr (int, optional): Threshold for NCR voxel count. Defaults to 50.
327        t_ed (int, optional): Threshold for ED voxel count. Defaults to 50.
328        t_et (int, optional): Threshold for ET voxel count. Defaults to 50.
329
330    Returns:
331        Tuple[int, int, int, int, int, int]: 
332            Number of removed NCR, total NCR, removed ED, total ED, removed ET, total ET.
333    """
334    seg_obj = nib.load(seg_file)
335    labels_ncr, labels_ed, labels_et, n_ncr, n_ed, n_et = get_connected_labels(seg_file)
336    
337    # Process NCR
338    vols = []
339    for i in range(n_ncr):
340        tmp = np.where(labels_ncr == i + 1, 1, 0)
341        vol = np.count_nonzero(tmp)
342        if vol < t_ncr:
343            labels_ncr = np.where(labels_ncr == i + 1, 0, labels_ncr)
344            vols.append(vol)
345    removed_ncr = len(vols)
346    
347    # Process ED
348    vols = []
349    for i in range(n_ed):
350        tmp = np.where(labels_ed == i + 1, 1, 0)
351        vol = np.count_nonzero(tmp)
352        if vol < t_ed:
353            labels_ed = np.where(labels_ed == i + 1, 0, labels_ed)
354            vols.append(vol)
355    removed_ed = len(vols)
356    
357    # Process ET
358    vols = []
359    for i in range(n_et):
360        tmp = np.where(labels_et == i + 1, 1, 0)
361        vol = np.count_nonzero(tmp)
362        if vol < t_et:
363            labels_et = np.where(labels_et == i + 1, 0, labels_et)
364            vols.append(vol)
365    removed_et = len(vols)
366
367    # Combine the cleaned labels
368    new_ncr = np.where(labels_ncr != 0, 1, 0)
369    new_ed = np.where(labels_ed != 0, 2, 0)
370    new_et = np.where(labels_et != 0, 3, 0)
371    new_seg = new_ncr + new_ed + new_et
372    new_obj = nib.Nifti1Image(new_seg.astype(np.int8), seg_obj.affine)
373    nib.save(new_obj, out_file)
374    return removed_ncr, n_ncr, removed_ed, n_ed, removed_et, n_et

Removes disconnected small components from the segmentation based on thresholds.

Args: seg_file (Path): Path to the input segmentation file. out_file (Path): Path to save the cleaned segmentation file. t_ncr (int, optional): Threshold for NCR voxel count. Defaults to 50. t_ed (int, optional): Threshold for ED voxel count. Defaults to 50. t_et (int, optional): Threshold for ET voxel count. Defaults to 50.

Returns: Tuple[int, int, int, int, int, int]: Number of removed NCR, total NCR, removed ED, total ED, removed ET, total ET.

def remove_disconnected_from_dir( input_dir: pathlib.Path, output_dir: pathlib.Path, t_ncr: int = 50, t_ed: int = 50, t_et: int = 50) -> pathlib.Path: 
377def remove_disconnected_from_dir(
378    input_dir: Path,
379    output_dir: Path,
380    t_ncr: int = 50,
381    t_ed: int = 50,
382    t_et: int = 50
383) -> Path:
384    """
385    Removes disconnected small components from all segmentation files in a directory.
386
387    Args:
388        input_dir (Path): Path to the input directory containing segmentation files.
389        output_dir (Path): Path to the output directory to save cleaned segmentation files.
390        t_ncr (int, optional): Threshold for NCR voxel count. Defaults to 50.
391        t_ed (int, optional): Threshold for ED voxel count. Defaults to 50.
392        t_et (int, optional): Threshold for ET voxel count. Defaults to 50.
393
394    Returns:
395        Path: Path to the output directory.
396    """
397    if not output_dir.exists():
398        output_dir.mkdir(parents=True)
399
400    for seg1 in input_dir.iterdir():
401        if seg1.name.endswith('.nii.gz'):
402            casename = seg1.name
403            seg2 = output_dir / casename
404            removed_ncr, n_ncr, removed_ed, n_ed, removed_et, n_et = remove_disconnected(
405                seg1, seg2, t_ncr=t_ncr, t_ed=t_ed, t_et=t_et
406            )
407            print(
408                f"{casename} removed regions NCR {removed_ncr:03d}/{n_ncr:03d} "
409                f"ED {removed_ed:03d}/{n_ed:03d} ET {removed_et:03d}/{n_et:03d}"
410            )
411        else:
412            print("Wrong input file!")
413
414    return output_dir

Removes disconnected small components from all segmentation files in a directory.

Args: input_dir (Path): Path to the input directory containing segmentation files. output_dir (Path): Path to the output directory to save cleaned segmentation files. t_ncr (int, optional): Threshold for NCR voxel count. Defaults to 50. t_ed (int, optional): Threshold for ED voxel count. Defaults to 50. t_et (int, optional): Threshold for ET voxel count. Defaults to 50.

Returns: Path: Path to the output directory.