Whole Slide Image Registration

Click to open in: [GitHub][Colab]

Importing related libraries

We import some standard Python modules, and also the TIAToolbox Python modules for the image registration task, written by the TIA Centre team.

from tiatoolbox.models.engine.semantic_segmentor import (
    IOSegmentorConfig,
    SemanticSegmentor,
)
from tiatoolbox.tools.registration.wsi_registration import match_histograms, DFBRegister
from tiatoolbox.utils.misc import imread
import os
import cv2
import numpy as np
from scipy import ndimage
import SimpleITK as sitk
from matplotlib import pyplot as plt
from skimage import color, exposure, measure, morphology
import matplotlib as mpl

mpl.rcParams["figure.dpi"] = 160  # for high resolution figure in notebook
mpl.rcParams["figure.facecolor"] = "white"  # To make sure text is visible in dark mode

Downloading the required files

We download, over the internet, image files used for the purpose of this notebook. In particular, we use a sample image pair from the ANHIR challenge dataset, which have been downloaded from https://anhir.grand-challenge.org/Data/. The challenge dataset is available under the following licence: CC-BY-NC-SA. The challenge paper [1] is available here.

[1] Borovec, Jiří, et al. “ANHIR: automatic non-rigid histological image registration challenge.” IEEE transactions on medical imaging 39.10 (2020): 3042-3052.

Downloading is needed once in each Colab session and it should take less than 1 minute.

In Colab, if you click the file’s icon (see below) in the vertical toolbar on the left-hand side then you can see all the files which the code in this notebook can access. The data will appear here when it is downloaded.

import requests

fixed_img_file_name = "fixed_image.jpg"
moving_img_file_name = "moving_image.jpg"

# Downloading fixed image from ANHIR challenge
r = requests.get(
    "https://tiatoolbox.dcs.warwick.ac.uk/testdata/registration/ANHIR/COAD_06_HE.jpg"
)
with open(fixed_img_file_name, "wb") as f:
    f.write(r.content)

# Downloading moving image from ANHIR challenge
r = requests.get(
    "https://tiatoolbox.dcs.warwick.ac.uk/testdata/registration/ANHIR/COAD_06_S2.jpg"
)
with open(moving_img_file_name, "wb") as f:
    f.write(r.content)

Reading and visualising images

fixed_image_rgb = imread(fixed_img_file_name)
moving_image_rgb = imread(moving_img_file_name)

_, axs = plt.subplots(1, 2, figsize=(15, 10))
axs[0].imshow(fixed_image_rgb, cmap="gray")
axs[0].set_title("Fixed Image")
axs[1].imshow(moving_image_rgb, cmap="gray")
axs[1].set_title("Moving Image")
plt.show()

Image Pre-processing

The images are converted to greyscale. To unify the appearance of an image pair, histogram matching is performed as a normalisation step.

def preprocess_image(image):
    image = color.rgb2gray(image)
    image = exposure.rescale_intensity(
        image, in_range=tuple(np.percentile(image, (0.5, 99.5)))
    )
    image = image * 255
    return image.astype(np.uint8)


fixed_image = preprocess_image(fixed_image_rgb)
moving_image = preprocess_image(moving_image_rgb)
fixed_image, moving_image = match_histograms(fixed_image, moving_image)

_, axs = plt.subplots(1, 2, figsize=(15, 10))
axs[0].imshow(fixed_image, cmap="gray")
axs[0].set_title("Fixed Image")
axs[1].imshow(moving_image, cmap="gray")
axs[1].set_title("Moving Image")
plt.show()

temp = np.repeat(np.expand_dims(fixed_image, axis=2), 3, axis=2)
cv2.imwrite(os.path.join(global_save_dir, "fixed.png"), temp)
temp = np.repeat(np.expand_dims(moving_image, axis=2), 3, axis=2)
cv2.imwrite(os.path.join(global_save_dir, "moving.png"), temp)

True

Tissue Segmentation

In this section, image segmentation is performed to generate tissue masks so that registration can be performed using only the active or discriminatory tissue area. In DFBR, these are used to exclude the matching points from the non-tissue and fatty regions while in the non-rigid alignment method, the background region is excluded from images before inputting them to the SimpleITK module.

We should mention that when you use a TIAToolbox pretrained model, you don’t need to worry about setting the input/output shape parameters as their optimal values will be loaded by default.

save_dir = os.path.join(global_save_dir, "tissue_mask")
if os.path.exists(save_dir):
    shutil.rmtree(save_dir, ignore_errors=False, onerror=None)

segmentor = SemanticSegmentor(
    pretrained_model="unet_tissue_mask_tsef",
    num_loader_workers=4,
    batch_size=4,
)

output = segmentor.predict(
    [
        os.path.join(global_save_dir, "fixed.png"),
        os.path.join(global_save_dir, "moving.png"),
    ],
    save_dir=save_dir,
    mode="tile",
    resolution=1.0,
    units="baseline",
    patch_input_shape=[1024, 1024],
    patch_output_shape=[512, 512],
    stride_shape=[512, 512],
    on_gpu=ON_GPU,
    crash_on_exception=True,
)

Post-processing and Visualization of Masks

In the output, the prediction method returns a list of the paths to its inputs and to the processed outputs saved on the disk. This can be used to load the results for processing and visualisation. In this section, a simple processing of the raw prediction is performed to generate the tissue mask.

def post_processing_mask(mask):
    mask = ndimage.binary_fill_holes(mask, structure=np.ones((3, 3))).astype(int)

    # remove all the objects while keep the biggest object only
    label_img = measure.label(mask)
    if len(np.unique(label_img)) > 2:
        regions = measure.regionprops(label_img)
        mask = mask.astype(bool)
        all_area = [i.area for i in regions]
        second_max = max([i for i in all_area if i != max(all_area)])
        mask = morphology.remove_small_objects(mask, min_size=second_max + 1)
    return mask.astype(np.uint8)


fixed_mask = np.load(output[0][1] + ".raw.0.npy")
moving_mask = np.load(output[1][1] + ".raw.0.npy")

# Simple processing of the raw prediction to generate semantic segmentation task
fixed_mask = np.argmax(fixed_mask, axis=-1) == 2
moving_mask = np.argmax(moving_mask, axis=-1) == 2

fixed_mask = post_processing_mask(fixed_mask)
moving_mask = post_processing_mask(moving_mask)

_, axs = plt.subplots(1, 2, figsize=(15, 10))
axs[0].imshow(fixed_mask, cmap="gray")
axs[0].set_title("Fixed Mask")
axs[1].imshow(moving_mask, cmap="gray")
axs[1].set_title("Moving Mask")
plt.show()

Registration using DFBR

In this section, we apply the DFBR method for aligning the moving image with respect to the fixed image.

dfbr_fixed_image = np.repeat(np.expand_dims(fixed_image, axis=2), 3, axis=2)
dfbr_moving_image = np.repeat(np.expand_dims(moving_image, axis=2), 3, axis=2)

df = DFBRegister()
dfbr_transform = df.register(
    dfbr_fixed_image, dfbr_moving_image, fixed_mask, moving_mask
)

# Visualization
original_moving = cv2.warpAffine(
    moving_image, np.eye(2, 3), fixed_image.shape[:2][::-1]
)
dfbr_registered_image = cv2.warpAffine(
    moving_image, dfbr_transform[0:-1], fixed_image.shape[:2][::-1]
)
dfbr_registered_mask = cv2.warpAffine(
    moving_mask, dfbr_transform[0:-1], fixed_image.shape[:2][::-1]
)

before_overlay = np.dstack((original_moving, fixed_image, original_moving))
dfbr_overlay = np.dstack((dfbr_registered_image, fixed_image, dfbr_registered_image))

_, axs = plt.subplots(1, 2, figsize=(15, 10))
axs[0].imshow(before_overlay, cmap="gray")
axs[0].set_title("Overlay Before Registration")
axs[1].imshow(dfbr_overlay, cmap="gray")
axs[1].set_title("Overlay After DFBR")
plt.show()

BSpline Transform

In this section, we perform registration of fixed and moving images using a multi-resolution BSpline method. The moving image here is a registered image after the application of the DFBR transform. For the given image pair, this method is shown to perform better with the inverted intensity values followed by background removal.

The BSpline approach involves a number of parameters that the user can experiment with. Changing these parameters would result in varying complexity and registration completion times. Feel free to play around with the non-rigid alignment parameters and to experiment with new images.

# Inverting intensity values
fixed_image_inv = 255 - fixed_image
moving_image_inv = 255 - dfbr_registered_image

# Background Removal
fixed_mask = np.array(fixed_mask != 0, dtype=np.uint8)
dfbr_registered_mask = np.array(dfbr_registered_mask != 0, dtype=np.uint8)
fixed_image_inv = cv2.bitwise_and(fixed_image_inv, fixed_image_inv, mask=fixed_mask)
moving_image_inv = cv2.bitwise_and(
    moving_image_inv, moving_image_inv, mask=dfbr_registered_mask
)

# Getting SimpleITK Images from numpy arrays
fixed_image_inv_sitk = sitk.GetImageFromArray(fixed_image_inv)
moving_image_inv_sitk = sitk.GetImageFromArray(moving_image_inv)
fixed_image_inv_sitk = sitk.Cast(fixed_image_inv_sitk, sitk.sitkFloat32)
moving_image_inv_sitk = sitk.Cast(moving_image_inv_sitk, sitk.sitkFloat32)

# Determine the number of BSpline control points
mesh_size = [3] * fixed_image_inv_sitk.GetDimension()
tx = sitk.BSplineTransformInitializer(
    image1=fixed_image_inv_sitk, transformDomainMeshSize=mesh_size
)
print("Initial Number of Parameters: {0}".format(tx.GetNumberOfParameters()))

R = sitk.ImageRegistrationMethod()
R.SetInitialTransformAsBSpline(tx, inPlace=True, scaleFactors=[1, 2, 5])
R.SetMetricAsMattesMutualInformation(50)
R.SetMetricSamplingStrategy(R.RANDOM)
R.SetMetricSamplingPercentage(0.2)

R.SetShrinkFactorsPerLevel([4, 2, 1])
R.SetSmoothingSigmasPerLevel([4, 2, 1])
R.SetOptimizerAsGradientDescentLineSearch(
    0.5, 100, convergenceMinimumValue=1e-4, convergenceWindowSize=5
)
R.SetInterpolator(sitk.sitkLinear)
outTx = R.Execute(fixed_image_inv_sitk, moving_image_inv_sitk)

Visualizing BSpline Results

resampler = sitk.ResampleImageFilter()
resampler.SetReferenceImage(fixed_image_inv_sitk)
resampler.SetInterpolator(sitk.sitkLinear)
resampler.SetDefaultPixelValue(1)
resampler.SetTransform(outTx)

moving_image_sitk = sitk.GetImageFromArray(dfbr_registered_image)
sitk_registered_moving_image_sitk = resampler.Execute(moving_image_sitk)
sitk_registered_moving_image = sitk.GetArrayFromImage(sitk_registered_moving_image_sitk)

bspline_overlay = np.dstack(
    (sitk_registered_moving_image, fixed_image, sitk_registered_moving_image)
)

_, axs = plt.subplots(1, 2, figsize=(15, 10))
axs[0].imshow(dfbr_overlay)
axs[0].set_title("DFBR Overlay")
axs[1].imshow(bspline_overlay)
axs[1].set_title("Bspline Overlay")
plt.show()