Homologous point transformer for multi-modality prostate image registration

Background

Previous attempts to automate the prostate image registration process through the use of iterative computational methods and deep learning have resulted in moderate success (Sood et al., 2021; Shao et al., 2021b). The recent iterative method RAPSODI reconstructs the three-dimensional prostate volume from two-dimensional slices (Rusu et al., 2020). To find reasonable initial alignments for each slide, the slices are registered to their adjacent histology slices. Then, RAPSODI iteratively estimates optimal affine and nonlinear transformations using gradient descent with sum of squared differences and Mattes mutual information as loss functions. While this method is robust, it is computationally expensive and requires careful manual adjustment of several hyperparameters. ProsRegNet, a deep learning-based alternative to RAPSODI, directly estimates parameters for affine and thin plate spline transforms using a neural network (Shao et al., 2021a). The design is largely patterned after CNNGeometric (Rocco, Arandjelovic & Sivic, 2017), a general image registration pipeline, and uses ResNet-101 (He et al., 2016), a pre-trained feature extraction convolutional neural network (CNN), to create feature maps from both histology and MR images. A custom correlation layer is used to capture the relationships between the sectors of these feature maps before a second convolutional neural network directly predicts the parameters of ProsRegNet’s transforms. While ProsRegNet significantly improves upon RAPSODI’s inference time, the resolution of the feature maps it uses is limited due to the correlation layer’s high memory consumption. Furthermore, this network structure is optimized for aligning the masks of the MR and histology images at the cost of intra-prostate alignment. In 2021, Shao et al. (2021b) applied a weakly-supervised deep learning approach to improve on the results of ProsRegNet. The network they developed used two separate convolutional input heads that are concatenated and followed by dense layers to estimate affine and nonlinear transformable parameters. It was trained with both masked and unmasked MR images and beat ProsRegNet in all common registration metrics. Despite this, the network is still fundamentally limited by the precision of a 4 × 4 thin plate spline transform.

Study goals

In this study, we propose a new cross-modality registration pipeline that relies on a deep learning network to estimate homologous points on histology and MR images. After the network predicts homologous points, the proposed pipeline uses the estimated points to fit a thin plate spline transform and warp the histology into MR space. This pipeline provides three novel advantages over previously described approaches. First, this pipeline is transform-invariant; any image transform that can be resolved from two point clouds may be used in registration. Second, while most automated registrations pipelines require masked MR images, this approach can register completely unmasked data. Finally, the homologous point pipeline allows for more flexibility in registration fine-tuning through point cloud post-processing than the previous approaches while maintaining the inference time advantage of deep learning.

Overview

Our registration pipeline consists of three major components: a histology landmark generation algorithm, a homologous point prediction network, and a transformation resolver (Fig. 1). The landmark generation algorithm selects a set of points on the input histology image using scale invariant feature transform (SIFT) (Lowe, 2004) and enforcing a tunable point density. The homologous point prediction network uses three input heads that accept the original histology image (512 × 512), corresponding T2-weighted MR image (512 × 512), and the generated landmarks from step 1 to solve for the homologous landmarks in MR space. Finally, the set of histology-space input landmarks generated by the landmark generation algorithm and the paired set of MR-space homologous points predicted by the deep learning model are used by the transformation resolver to calculate a nonlinear transform (thin plate spline) that will warp the histology image to align with the MR image. This thin plate spline transform is finally applied to the histology image. The same transform may be applied to annotation maps derived in histology space to transfer pathological labels into MR space. All code and data used for this registration pipeline may be found at https://github.com/9crusher/homologous_point_prediction.

Data collection

The data set for this project was obtained from the Medical College of Wisconsin. All patients were recruited prospectively with institutional review board approval (Medical College of Wisconsin Internal Review Board—PRO00022426) and written consent. Patients were scheduled for radical prostatectomy with clinical MRI scans acquired approximately 2 weeks prior to surgery.

We follow the procedure described by McGarry et al. (2018); following prostatectomy, the prostate tissue was formalin fixed, and a custom slicing jig derived from the MRI scans was created to obtain histology/MR slice correspondence. Whole mount prostate slides were stained with H&E (hematoxylin and eosin) and imaged on a microscope slide scanner (captured using 40× objective at 0.58 $\mu$m/px). Manually-curated homologous control points were placed for all slides in the selected data set to act as ground truth. While full field MR images were used as input and training data for our model, prostate masks were also generated for use in calculating Dice coefficient (Dice, 1945; Sorenson, 1948) post training.

A total of 232 individual prostate slides from 45 patients were included in the study. Two patients were reserved for validation and seven were held-out for testing. The remaining patients comprised the training data set. Several slides included in the data set exhibit tears and/or folds. These artifacts were introduced during the slicing, staining, and/or mounting process.

Landmark selection algorithm

The selection algorithm was used to generate a set of points (landmarks) with uniform coverage of the input histology images. To improve the chances of a network finding the same points in both the MR and histology image, the algorithm was designed to select “high information” points on the histology image. These points lie along the edges of the prostate, along the boundaries of the urethra and seminal vesicles, and along the edges of other significant structural features. The landmark selection algorithm is capable of producing point sets of varying density based on a configuration parameter.

The algorithm first selects candidate output points using the same difference of Gaussian-based method described in the original SIFT paper by Lowe (2004). Candidate points with contrast scores below an empirically-derived threshold of 0.04 are eliminated. This threshold was chosen because it led to sets of points that closely matched the point density seen in the manually-created point sets. Next, an edge score is calculated for each remaining candidate point. Edge scores are calculated as the ratio of the larger to the smaller eigenvalue of the Hessian matrix at the location of the candidate point in scale space. Rather than directly calculating these values, the algorithm uses the computationally-efficient method described in Harris & Stephens (1988) to estimate edge scores. Candidate points in the empirically-chosen top 19th percentile of edge score are classified as edges, and the remaining points are classified as interior feature points. Finally, a minimum distance between points is enforced through the sequential elimination of points within some user-defined distance of other points. Edge points are prioritized in this elimination process so interior points are deleted in the case of an edge and interior point minimum distance conflict. The output of this algorithm is a list of landmark points with uniform density and and complete coverage of the histological surface.

Homologous point network

Network input

The network predicts the homologous points in MR space for some set of input landmark coordinates $\mathbf{X}\in {\mathbb{R}}^{L\times 2}$ where L is the number of points. It contains three input heads: one for the input point coordinates X, one for the moving image $\mathbf{M}\in {\mathbb{R}}^{512\times 512}$ and one for the fixed image $\mathbf{F}\in {\mathbb{R}}^{512\times 512}$. The output is a set of homologous point coordinates $\mathbf{Y}\in {\mathbb{R}}^{L\times 2}$ in MR space (Fig. 2A). Because L must be fixed for each input in a batch and point sets generated by the SIFT algorithm may have varying numbers of points, L is set to a number larger than the maximum expected number of points, and X is padded with [0, 0] points to achieve length L for every input. In practice, we found L = 75 to be large enough to contain all points generated by the point placement algorithm.

Model training

The homologous point prediction model was trained using the ground truths of manually-placed homologous points in MR and histology space. Mean absolute error between the predicted points and the ground truth points was used as the loss function. The network was trained for 128 epochs with the data augmentation described below. The model contains 5,371,904 trainable parameters and was run on a NVIDIA Tesla T4 GPU using the Tensorflow and Keras Python frameworks.

Evaluation metrics

Average control point deviation and Dice coefficient were used to quantitatively evaluate the registration performance. To calculate average control point deviation, a registration pipeline was used to solve for a transform from histology space to MR space. The human-placed landmarks in histology space were transformed to MR space using the solved transform. Finally, the average Euclidean distance between every pair of warped landmarks and ground truth MR points was calculated. Because the ground truth control points were placed on both the edges and interior of the prostate structure, this metric quantifies the edge and internal structural alignment of a registration result.

For clarification, when evaluating the homologous point pipeline with the control point deviation metric, two distinct point sets were at play. (1) The set of points placed by the SIFT algorithm and predicted in MR space by the deep learning network and (2) the human annotated landmarks in histology space and their corresponding human annotated landmarks in MR space. The quality of the resolved transform was evaluated using the human-placed ground truths.

Dice coefficient was also used to evaluate model performance. Visualizations of the evaluation metrics are shown in Fig. 3.

Softmax output—point placement visualization

To provide a better qualitative interpretation of how the network resolves the homologous points, we visualized the softmax layer output for a single homologous point in histology space—Fig. 4 (left). Figure 4 (center) depicts a single row of S projected back into the two dimensional space of F. To calculate a row in Y, the network resolves the centroid of all patches in the image modified by the depicted weights. This centroid is then further shifted by the original offset of the landmark in the histology side patch. This results in the predicted location of the homologous point in the fixed image—Fig. 4.

Method comparison

The proposed homologous point pipeline was compared to thin plate spline transforms resolved with human ground truth points and to several versions of the ProsRegNet, the current state-of-the-art model. The highest-performing published weights for ProsRegNet were used in the evaluation (ProsRegNet—Paper). Additionally, ProsRegNet was configured to train on the same supercomputer as the homologous point pipeline and was finetuned with a hyperparameter grid search on the same training data set. The highest performing finetuned ProsRegNet model (ProsRegNet—Tuned) was selected for this comparison. This step was performed to maximize the performance of ProsRegNet and allow it to generalize on the same cohort as the homologous point pipeline. All four models were evaluated for both Dice coefficient and mean control point deviation on the same held-out test set.

The average Dice coefficient for the homologous point pipeline was found to be $0.90\pm 0.05$, which under-performed Prosregnet—Paper at $0.94\pm 0.03$ $(p<0.001)$, Prosregnet—Tuned at $0.94\pm 0.03$ $(p<0.001)$, and the Human GT at $0.92\pm 0.05$ $(p=0.012)$ (Fig. 5—Top).

Mean control point deviation in the homologous point pipeline was found to be $9.55\pm 1.76$ pixels which outperformed both ProsRegNet—Paper at $13.44\pm 3.71$ pixels $(p<0.001)$ and ProsRegNet—Tuned at $13.12\pm 3.49$ pixels $(p<0.001)$ but fell short of the Human GT at $0.81\pm 0.16$ pixels $(p<0.001)$ (Fig. 5—Bottom).

It should be noted that the ProsRegNet pipeline requires masked MR images to resolve transforms while the homologous point pipeline registered unmasked images. Additionally, the control point deviation of the human ground truth registrations was near zero because this registration method directly used the ground truths to resolve TPS transforms which due to degrees of freedom may not perfectly fit all control points.

Future directions

With our homologous point pipeline established, we envision three broad categories of future work. First, the pipeline should be characterized on frequently seen histological artifacts; we have anecdotally observed encouraging performance with partial tissue samples and tears in the current data set. However, the performance on slides with several different types of artifacts including partial tissue sections, tearing, and folding should be quantified. These artifacts are commonly introduced during tissue processing and may exclude a significant number of samples from analysis. Second, this approach should be fine-tuned and tested on additional organ systems. In this study, the homologous point pipeline was only characterized for the prostate data, but it is designed to adapt easily to other data. Adaption of this pipeline to the brain is especially enticing due the size of the organ often limiting histological processing to partial tissue sections. To the best of our knowledge, other current approaches can not be adapted to resolve points between partial tissue samples and whole organ imaging. Finally, this current approach requires the given histology and MR slices to be paired at the input. Both the histology and MR samples may include three-dimensional relationships that this approach may not currently leverage. If this approach were adapted to organ volumes/voxels, we may improve homologous point predictions and eliminate the required pre-processing of tissue to MR slice pairing.