from monai.transforms import Transform, Compose, LoadImage, EnsureChannelFirst
import torch
import skimage
import torch
import SimpleITK as sitk
import numpy as np
from PIL import Image
from io import BytesIO
import matplotlib.pyplot as plt
import SimpleITK as sitk
from matplotlib.colors import ListedColormap
import base64
import numpy as np
from cv2 import dilate
from scipy.ndimage import label
from Model_Seg import RgbaToGrayscale
def image_to_base64(image_path):
with open(image_path, "rb") as image_file:
return base64.b64encode(image_file.read()).decode('utf-8')
class CustomCLAHE(Transform):
"""Implements Contrast-Limited Adaptive Histogram Equalization (CLAHE) as a custom transform, as described by Qiu et al.
Attributes:
p1 (float): Weighting factor, determines degree of of contour enhacement. Default is 0.6.
p2 (None or int): Kernel size for adaptive histogram. Default is None.
p3 (float): Clip limit for histogram equalization. Default is 0.01.
"""
def __init__(self, p1=0.6, p2=None, p3=0.01):
self.p1 = p1
self.p2 = p2
self.p3 = p3
def __call__(self, data):
"""Apply the CLAHE algorithm to input data.
Args:
data (Union[dict, np.ndarray]): Input data. Could be a dictionary containing the image or the image array itself.
Returns:
torch.Tensor: Transformed data.
"""
if isinstance(data, dict):
im = data["image"]
else:
im = data
im = im.numpy()
# remove the first dimension
im = im[0]
im = im[None, :, :]
#im = np.expand_dims(im, axis=0)
im = skimage.exposure.rescale_intensity(im, in_range="image", out_range=(0, 1))
im_noi = skimage.filters.median(im)
im_fil = im_noi - self.p1 * skimage.filters.gaussian(im_noi, sigma=1)
im_fil = skimage.exposure.rescale_intensity(im_fil, in_range="image", out_range=(0, 1))
im_ce = skimage.exposure.equalize_adapthist(im_fil, kernel_size=self.p2, clip_limit=self.p3)
if isinstance(data, dict):
data["image"] = torch.Tensor(im_ce)
else:
data = torch.Tensor(im_ce)
return data
def custom_colormap():
cdict = [(0, 0, 0, 0), # Class 0 - fully transparent (background)
(0, 1, 0, 0.5), # Class 1 - Green with 50% transparency
(1, 0, 0, 0.5), # Class 2 - Red with 50% transparency
(1, 1, 0, 0.5)] # Class 3 - Yellow with 50% transparency
cmap = ListedColormap(cdict)
return cmap
def read_image(image_path):
read_transforms = Compose([
LoadImage(image_only=True),
EnsureChannelFirst(),
RgbaToGrayscale(), # Convert RGBA to grayscale
])
try:
original_image = read_transforms(image_path)
original_image_np = original_image.numpy().astype(np.uint8)
return original_image_np.squeeze()
except Exception as e:
try :
original_image = sitk.ReadImage(image_path)
original_image_np = sitk.GetArrayFromImage(original_image)
return original_image_np.squeeze()
except Exception as e:
print("Failed Loading the Image: ", e)
return None
def overlay_mask(image_path, image_mask):
original_image_np = read_image(image_path).squeeze().astype(np.uint8)
#adjust mask intensities for display
image_mask_disp = image_mask
plt.figure(figsize=(10, 10))
plt.imshow(original_image_np, cmap='gray')
plt.imshow(image_mask_disp, cmap=custom_colormap(), alpha=0.5)
plt.axis('off')
# Save the overlay to a buffer
buffer = BytesIO()
plt.savefig(buffer, format='png', bbox_inches='tight', pad_inches=0)
buffer.seek(0)
overlay_image_np = np.array(Image.open(buffer))
return overlay_image_np, original_image_np
def bounding_box_mask(image, label):
"""Generates a bounding box mask around a labeled region in an image
Args:
image (SimpleITK.Image): The input image.
label (SimpleITK.Image): The labeled image containing the region of interest.
Returns:
SimpleITK.Image: An image containing the with the bounding box mask applied with the
same spacing as the original image.
Note:
This function assumes that the input image and label are SimpleITK.Image objects.
The returned bounding box mask is a binary image where pixels inside the bounding box
are set to 1 and others are set to 0.
"""
# get original spacing
original_spacing = image.GetSpacing()
# convert image and label to arrays
image_array = sitk.GetArrayFromImage(image)
image_array = np.squeeze(image_array)
label_array = sitk.GetArrayFromImage(label)
label_array = np.squeeze(label_array)
# determine corners of the bounding box
first_nonzero_row_index = np.nonzero(np.any(label_array != 0, axis=1))[0][0]
last_nonzero_row_index = np.max(np.nonzero(np.any(label_array != 0, axis=1)))
first_nonzero_column_index = np.nonzero(np.any(label_array != 0, axis=0))[0][0]
last_nonzero_column_index = np.max(np.nonzero(np.any(label_array != 0, axis=0)))
top_left_corner = (first_nonzero_row_index, first_nonzero_column_index)
bottom_right_corner = (last_nonzero_row_index, last_nonzero_column_index)
# define the bounding box as an array mask
bounding_box_array = label_array.copy()
bounding_box_array[
top_left_corner[0] : bottom_right_corner[0] + 1,
top_left_corner[1] : bottom_right_corner[1] + 1,
] = 1
# add channel dimension
bounding_box_array = bounding_box_array[None, ...].astype(np.uint8)
# get Image from Array Mask and apply original spacing
bounding_box_image = sitk.GetImageFromArray(bounding_box_array)
bounding_box_image.SetSpacing(original_spacing)
return bounding_box_image
def threshold_based_crop(image):
"""
Use Otsu's threshold estimator to separate background and foreground. In medical imaging the background is
usually air. Then crop the image using the foreground's axis aligned bounding box.
Args:
image (SimpleITK image): An image where the anatomy and background intensities form a
bi-modal distribution
(the assumption underlying Otsu's method.)
Return:
Cropped image based on foreground's axis aligned bounding box.
"""
inside_value = 0
outside_value = 255
label_shape_filter = sitk.LabelShapeStatisticsImageFilter()
# uncomment for debugging
#sitk.WriteImage(image, "./image.png")
label_shape_filter.Execute(sitk.OtsuThreshold(image, inside_value, outside_value))
bounding_box = label_shape_filter.GetBoundingBox(outside_value)
return sitk.RegionOfInterest(
image,
bounding_box[int(len(bounding_box) / 2) :],
bounding_box[0 : int(len(bounding_box) / 2)],
)
def creat_SIJ_mask(image, input_label):
"""
Create a mask for the sacroiliac joints (SIJ) from pelvis and sascrum segmentation mask
Args:
image (SimpleITK.Image): x-ray image.
input_label (SimpleITK.Image): Segmentation mask containing labels for sacrum, left- and right pelvis
Returns:
SimpleITK.Image: Mask of the SIJ
"""
original_spacing = image.GetSpacing()
# uncomment for debugging
#sitk.WriteImage(input_label, "./input_label.png")
mask_array = sitk.GetArrayFromImage(input_label).squeeze()
sacrum_value = 1
left_pelvis_value = 2
right_pelvis_value = 3
background_value = 0
sacrum_mask = (mask_array == sacrum_value)
first_nonzero_column_index = np.nonzero(np.any(sacrum_mask != 0, axis=0))[0][0]
last_nonzero_column_index = np.max(np.nonzero(np.any(sacrum_mask != 0, axis=0)))
box_width=last_nonzero_column_index-first_nonzero_column_index
dilation_extent = int(np.round(0.05 * box_width))
dilated_sacrum_mask = dilate_mask(sacrum_mask, dilation_extent)
intersection_left = (dilated_sacrum_mask & (mask_array == left_pelvis_value))
if np.all(intersection_left == 0):
print("Warning: No left intersection")
left_pelvis_mask = (mask_array == 2)
intersection_left = create_median_height_array(left_pelvis_mask)
intersection_left = keep_largest_component(intersection_left)
intersection_right = (dilated_sacrum_mask & (mask_array == right_pelvis_value))
if np.all(intersection_right == 0):
print("Warning: No right intersection")
right_pelvis_mask = (mask_array == 3)
intersection_right = create_median_height_array(right_pelvis_mask)
intersection_right = keep_largest_component(intersection_right)
intersection_mask = intersection_left +intersection_right
intersection_mask = intersection_mask[None, ...]
instersection_mask_im = sitk.GetImageFromArray(intersection_mask)
instersection_mask_im.SetSpacing(original_spacing)
return instersection_mask_im
def dilate_mask(mask, extent):
"""
Keeps only the largest connected component in a binary segmentation mask.
Args:
mask (numpy.ndarray): A numpy array representing the binary segmentation mask,
with 1s indicating the label and 0s indicating the background.
Returns:
numpy.ndarray: A modified version of the input mask, where only the largest
connected component is retained, and other components are set to 0.
"""
mask_uint8 = mask.astype(np.uint8)
kernel = np.ones((2*extent+1, 2*extent+1), np.uint8)
dilated_mask = dilate(mask_uint8, kernel, iterations=1)
return dilated_mask
def mask_and_crop(image, input_label):
"""
Performs masking and cropping operations on an image and its label.
Args:
image (SimpleITK.Image): The image to be processed.
label (SimpleITK.Image): The corresponding label image.
Returns:
tuple: A tuple containing two SimpleITK.Image objects.
- cropped_boxed_image: The image after applying bounding box masking and cropping.
- mask: The binary mask corresponding to the label after cropping.
Note:
This function relies on other functions: bounding_box_mask() and threshold_based_crop().
"""
input_label = creat_SIJ_mask(image,input_label)
box_mask = bounding_box_mask(image, input_label)
boxed_image = sitk.Mask(image, box_mask, maskingValue=0, outsideValue=0)
masked_image = sitk.Mask(image, input_label, maskingValue=0, outsideValue=0)
cropped_boxed_image = threshold_based_crop(boxed_image)
cropped_masked_image = threshold_based_crop(masked_image)
mask = np.squeeze(sitk.GetArrayFromImage(cropped_masked_image))
mask = np.where(mask > 0, 1, 0)
mask = sitk.GetImageFromArray(mask[None, ...])
return cropped_boxed_image, mask
def create_median_height_array(mask):
"""
Creates an array based on the median height of non-zero elements in each column of the input mask.
Args:
mask (numpy.ndarray): A binary mask with 1s representing the label and 0s the background.
Returns:
numpy.ndarray: A new binary mask array with columns filled based on the median height,
or None if the input mask has no non-zero columns.
Note:
This function is only used when there is no intersection between pelvis and sacrum, and creates an alternative
SIJ mask, that serves as an approximate replacement.
"""
rows, cols = mask.shape
column_details = []
for col in range(cols):
column_data = mask[:, col]
non_zero_indices = np.nonzero(column_data)[0]
if non_zero_indices.size > 0:
height = non_zero_indices[-1] - non_zero_indices[0] + 1
start_idx = non_zero_indices[0]
column_details.append((height, start_idx, col))
if not column_details:
return None
median_height = round(np.median([h[0] for h in column_details]))
median_cols = [(col, start_idx) for height, start_idx, col in column_details if height == median_height]
new_array = np.zeros_like(mask, dtype=int)
for col, start_idx in median_cols:
start_col = max(0, col - 5)
end_col = min(cols, col + 5)
new_array[start_idx:start_idx + median_height, start_col:end_col] = 1
return new_array
def keep_largest_component(mask):
"""
Identifies and retains the largest connected component in a binary segmentation mask.
Args:
mask (numpy.ndarray): A binary mask with 1s representing the label and 0s the background.
Returns:
numpy.ndarray: The modified mask with only the largest connected component.
"""
# Label the connected components
labeled_array, num_features = label(mask)
# If no features are found, return the original mask
if num_features <= 1:
return mask
# Find the largest connected component
largest_component = np.argmax(np.bincount(labeled_array.flat)[1:]) + 1
# Generate the mask for the largest component
return (labeled_array == largest_component).astype(mask.dtype)