Create utils.py

utils.py

@@ -0,0 +1,147 @@
+import numpy as np
+import SimpleITK as sitk
+
+channels = [
+    "background",
+    "spleen",
+    "right_kidney",
+    "left_kidney",
+    "gallbladder",
+    "liver",
+    "stomach",
+    "pancreas",
+    "right_adrenal_gland",
+    "left_adrenal_gland",
+    "left_lung",
+    "right_lung",
+    "heart",
+    "aorta",
+    "inferior_vena_cava",
+    "portal_vein_and_splenic_vein",
+    "left_iliac_artery",
+    "right_iliac_artery",
+    "left_iliac_vena",
+    "right_iliac_vena",
+    "esophagus",
+    "small_bowel",
+    "duodenum",
+    "colon",
+    "urinary_bladder",
+    "spine",
+    "sacrum",
+    "left_hip",
+    "right_hip",
+    "left_femur",
+    "right_femur",
+    "left_autochthonous_muscle",
+    "right_autochthonous_muscle",
+    "left_iliopsoas_muscle",
+    "right_iliopsoas_muscle",
+    "left_gluteus_maximus",
+    "right_gluteus_maximus",
+    "left_gluteus_medius",
+    "right_gluteus_medius",
+    "left_gluteus_minimus",
+    "right_gluteus_minimus",
+]
+
+
+def make_isotropic(image, interpolator=sitk.sitkLinear, spacing=None):
+    """
+    Many file formats (e.g. jpg, png,...) expect the pixels to be isotropic, same
+    spacing for all axes. Saving non-isotropic data in these formats will result in
+    distorted images. This function makes an image isotropic via resampling, if needed.
+    Args:
+        image (SimpleITK.Image): Input image.
+        interpolator: By default the function uses a linear interpolator. For
+                      label images one should use the sitkNearestNeighbor interpolator
+                      so as not to introduce non-existant labels.
+        spacing (float): Desired spacing. If none given then use the smallest spacing from
+                         the original image.
+    Returns:
+        SimpleITK.Image with isotropic spacing which occupies the same region in space as
+        the input image.
+    """
+    original_spacing = image.GetSpacing()
+    # Image is already isotropic, just return a copy.
+    if all(spc == original_spacing[0] for spc in original_spacing):
+        return sitk.Image(image)
+    # Make image isotropic via resampling.
+    original_size = image.GetSize()
+    if spacing is None:
+        spacing = min(original_spacing)
+    new_spacing = [spacing] * image.GetDimension()
+    new_size = [int(round(osz * ospc / spacing)) for osz, ospc in zip(original_size, original_spacing)]
+    return sitk.Resample(
+        image,
+        new_size,
+        sitk.Transform(),
+        interpolator,
+        image.GetOrigin(),
+        new_spacing,
+        image.GetDirection(),
+        0,  # default pixel value
+        image.GetPixelID(),
+    )
+
+
+def label_mapper(seg):
+
+    labels = []
+    for _class in np.unique(seg):
+        if _class == 0:
+            continue
+        labels.append((seg == _class, channels[_class]))
+
+    return labels
+
+
+def sitk2numpy(img, normalize=False):
+    img = sitk.DICOMOrient(img, "LPS")
+    # img = make_isotropic(img)
+    img = sitk.GetArrayFromImage(img)
+    if normalize:
+        minval, maxval = np.min(img), np.max(img)
+        img = ((img - minval) / (maxval - minval)).clip(0, 1) * 255
+    img = img.astype(np.uint8)
+    return img
+
+
+def read_image(path, normalize=False):
+
+    img = sitk.ReadImage(path)
+    return sitk2numpy(img, normalize)
+
+
+def display(image, seg=None, _slice=50):
+
+    # Image
+    if image is None or (isinstance(image, list) and len(image) == 0):
+        return None
+    if isinstance(image, list):
+        image = image[-1]
+    x = int(_slice * (image.shape[0] / 100))
+    image = image[x, :, :]
+
+    # Segmentation
+    if seg is None or (isinstance(seg, list) and len(seg) == 0):
+        seg = []
+    else:
+        if isinstance(seg, list):
+            seg = seg[-1]
+        seg = label_mapper(seg[x, :, :])
+
+    return image, seg
+
+
+def read_and_display(path, image_state, seg_state):
+
+    image_state.clear()
+    seg_state.clear()
+
+    if path is not None:
+        image = read_image(path, normalize=True)
+        image_state.append(image)
+        return display(image), image_state, seg_state
+    else:
+        return None, image_state, seg_state