Data Stream Stabilization for Optical Coherence Tomography Volumetric Scanning

The first part of the stabilization algorithm follows two steps to estimate the NURD vector between the previous stabilized ${{\hat{\boldsymbol{F}}}_{k-1}}$ and pre-shifted frame ${{\tilde{\boldsymbol{F}}}_{k}}$: First a correlation map ${{\boldsymbol{M}_{k}}}$ is calculated using cross-correlation, and then a CNN estimates the NURD vector taking the ${{\boldsymbol{M}_{k}}}$ as input.

Given two consecutive OCT frames with height $H$ (with $H$ A-lines in one frame) and width $W$ in polar coordinates, we used the Pearson correlation coefficient to reflect the similarity between two local image patches from these two frames:

where $\boldsymbol{f}_{i}\in{{\mathbb{R}}^{h\times W}}$ ($h\ll H$) is a local image patch centered at index position $i,i\in[0,H)$ of the newest frame. Correspondingly, ${{{\boldsymbol{f}}}_{j}^{\prime}}\in{{\mathbb{R}}^{h\times W}}$ is a local image patch centered at index position $j$ of ${{\hat{\boldsymbol{F}}}_{k-1}}$ ( $j\in[i-w/2,i+w/2)$, $w$ is is the shifting window width). ${s}_{i,j}$ is one element of the correlation map ${{\boldsymbol{M}_{k}}}\in{{\mathbb{R}}^{H\times w}}$. The pixel index $l$ will operate through the patch size $n=h\times W$. $\bar{\boldsymbol{f}_{i}}$ and $\bar{\boldsymbol{f}_{j}^{\prime}}$ are the mean values of patches ${\boldsymbol{f}_{i}}$ and ${\boldsymbol{f}_{j}^{\prime}}$ respectively.

The correlation map ${{\boldsymbol{M}_{k}}}$ contains information on the NURD, and we design a CNN to further extract the NURD vector. The detailed diagram of the NURD estimating nets is shown in fig. 4, where each block gives the convolutional kernel size, the stride step length and the output feature map depth of each convolution layer. Two convolution sub-branches are operated in parallel to capture features and produce hierarchically coarse-to-fine responses. The main difference between these two sub-branches is the strides control. Both the left sub-branch and the right sub-branch of the NURD estimation nets have 6 convolutional layers, and a LeakyReLU activation [49] is used after each convolution layer.

The left sub-branch always keeps the vertical stride as 1. The objective of this design is to emphasize the vertical spacial correspondence. By doing so, the front 5 feature extraction layers can gradually reduce the feature map width from $w$ to 1, while maintaining the feature map height $H$ as the input’s height. The depth of each convolution operation’s output is twice as deep as its input (here we set the output depth of the first layer as 8). The $5th$ feature map $\boldsymbol{A}_{F}^{5}\in{{\mathbb{R}}^{H\times 1\times 128}}$ extracts 128 local features, which could include the minimal value position, edge, and boundary position. A final layer with kernel size $1\times 1$ and channel depth 128, reorganizes the $5th$ feature map and decrease channels to a sub-branch output of $\bar{P}$ with size $H\times 1$.

Considering that in certain situations ${{\boldsymbol{M}_{k}}}$ will miss valid information for some row $m_{i}$ when there is no feature in a patch (window) $\boldsymbol{f}_{i}$ of ${{\tilde{\boldsymbol{F}}}_{k}}$, the right sub-branch is added. This sub-branch loosens the horizontal stride to 2 to involve more spacial information. This is a more common architecture for many deep learning applications [45]. Compared to the left sub-branch, this sub-branch can extract higher abstract features which are less sensitive to local missing information in the correlation map. We put larger kernels in the front layers to involve more common spatial information for neighboring positions of convolution, which prevents the under-fitting of feature extraction [16]. The kernel size is reduced in deeper layers to ensure a good trade-off between kernel size and computational burden. After the six convolutions, this subbranch outputs a matrix $\bar{\boldsymbol{P}}_{m}^{\prime\prime}$ of size $(H/2^{6})\times w$. To get an output of warping vector size $H\times 1$, a reshape operation is applied to the output of this sub-branch (i.e. by sequentially connecting each $1\times w$ row). A final convolution layer with kernel size $3\times 1$, input depth 2 and output depth 1, combines the prediction of two sub-branches to get an estimation of $\bar{P}$. The loss function for training the NURD estimating nets uses the conventional $L_{2}$ loss and an additional continuity loss. A standard $L_{2}$ loss is described by Eq. (2):

where $P_{i}$ is the element of $P$, the true NURD vector (ground truth), $n_{p}=H$ is the vector length. The $L_{2}$ loss function is commonly used for value estimation, while for this estimation task, to take into account the prior knowledge on the continuity of the distortion vector [8, 9, 10], we add the continuity loss as follows:

By calculating $L_{c}$, and combining it with $L_{2}$ in the network training, the attraction towards local minima with discontinuous NURD estimation will be suppressed. The final loss for branch (A) is:

It is possible to correct a short frame stack with just the NURD estimation part presented in the previous subsection. To compensate the accumulative error through iterations (especially for a long data stream), we estimate an overall rotation value ${{\bar{r}}_{k}}$.

In conventional pullback scans only the sheath data is consistent no matter what the environment outside the sheath is. According to this, the overall rotation can be observed by matching real-time sheath images with pre-recorded reference sheath images. However, when using the unstabilized pullback scanning to record the reference images, it still suffers from rotational distortion. To ensure that the orientations of reference frames are correct, we follow a calibration procedure. The setup of the sheath registration is shown in Fig. 5 (A), which relies on a external calibration object (a straight and flat ruler). As shown in Fig. 5 (B), the raw reference data still has the rotational distortion, which leads to both the ruler and sheath images located at wrong direction. We extract the contour surface of the ruler, and align the raw reference frame stack by minimizing the surface distance of all frames. By doing so, the rational error of the raw reference volumetric data is reduced from 59.4^∘ to 2.79^∘ (see Fig. 5 (C)). This calibrated reference data composes one of the inputs of the overall rotation estimation. Another input is composed by real-time B-scans, where image outside the sheath is masked out and only the sheath part is used.

With a recorded sheath image stack $\boldsymbol{S}\in{{\mathbb{R}}^{H\times W\times N}}$ (N is the number of frames in the entire reference stack), we compute Euler distances to estimate the rotation deviation $\Delta\bar{r}_{k}$ of pre-shifted image ${{\tilde{\boldsymbol{F}}}_{k}}$ (pre-shifted by $\bar{r}_{k-1}$):

where ${{\tilde{S}}_{k}}\in{{\mathbb{R}}^{H\times W\times l}}$ is buffered pre-shifted original images, $S_{k}\in{{\mathbb{R}}^{H\times W\times l}}$ is the corresponding selected buffer indexed at $k$ of $\boldsymbol{S}$ ($l\ll N$, and assuming that in every pullback scanning the torque coil moves with equal interval). The function $\mathcal{L}$ calculates Euler distance of two input buffers. For the pullback scan, $S_{k}$ is the corresponding reference image in the data stack, while to correct an outer robotic tool pullback 3D scan the algorithm can be simplified and just  use the first frame ${{\boldsymbol{F}}}_{0}$ as the reference, because in this scanning condition the relative longitude position between sheath and rotating ball-lens is constant. For both ${{\tilde{S}}_{k}}$ and $S_{k}$ only the sheath part is used for distance calculation. The index $r\in(-w,w)$ shifts within a given window $w$ to calculate the distance $d_{r}$, and the position with the lowest distance value is considered as the matched angular position of $\Delta\bar{r}_{k}$. Finally, the overall rotation used to compensate the NURD drift is computed by $\bar{r}_{k}=\bar{r}_{k-1}+\Delta\bar{r}_{k}$.

We use the concept of a PI Complementary Filter [31] to fuse the $\bar{P}_{k}$ vector with the $\bar{r}_{k}$ value, which is shown to be efficient and able to solve the warping vector integral drift problem. A discrete form of PI complementary filter is:

where $k_{p}$ and $k_{i}$ are PI compensating gains. ${{I}_{k}}$ is the integral component vector. $\boldsymbol{1}$ is a vector of ones. Each element ${{\hat{P}}_{k,i}}$ of $\hat{P_{k}}$ represents the angular distance between the position of the $i$th A-line of ${{\tilde{\boldsymbol{F}}}_{k}}$ and its correct position in polar domain. Applying $\hat{P_{k}}$ to this raw frame ${{\tilde{\boldsymbol{F}}}_{k}}$, a stabilized frame ${{\hat{\boldsymbol{F}}}_{k}}$ is obtained.

We test the proposed stabilization algorithm for both the internal conventional pullback and the outer scanning using a custom made OCT system (see description in section III.A). We collect videos with the testing object (fig. 2 (C)) to evaluate the performance for both types of scanning mode, since it has a special geometry. 6 scans were collected for each scanning mode, and each data stream comprises 500 frames. For the internal pullback, one reference scan is collected following the procedure described in section III.B.(2), and this reference scan is used in the correction of all internal pullback scans. To perform quantitative analysis, we also use synthetic videos generated from a variety of OCT images [50, 51, 52] to test the proposed algorithm, since it is almost impossible to manually annotate the NURD in the OCT data. We used the published videos from [50, 51, 52], and extracted all the OCT images out and transformed them into the same resolution in the polar domain. Approximately 5000 images which cover the catheter based OCT imaging in intravascular, digestive tract and lung respiration airway are obtained from these public videos.

The network is only trained with synthetic videos by intentionally adding NURD artifacts to the Optical Coherence Tomography (OCT) images, and tested on both true OCT scans and synthetic scanning videos. To ensure that the artificially added distortion covers the distribution of artifacts present in real videos, we first compute the correlation map $\boldsymbol{M}_{k}$ of every neighboring raw image pairs in real videos, and then measure the maximum NURD error with a graph path searching method [11]. The range of the GS-observed NURD was then expanded by 1/3 on each side. For each synthetic image, the synthetic NURD vector was then randomly generated within this expanded range. In addition, we implemented data augmentation that included geometric transformations, noise addition, brightness/contrast modification and OCT speckle/shadow simulation. By doing so the domain of artifacts of synthetic videos should cover the distribution of real artifacts. We randomly split all the collected images by 1:1:1 into train, validation and test data. For both the CNN training and algorithm deployment, we use the same PC with an Nvidia Qt1000 graphic card and Intel i5-9400H CPU. The code is implemented using the Pytorch framework[53]. For training, Batch Normalization (BN) is used right after convolution and before activation [54], dropout is not used [58] and weight initialization is performed following the method described in [57]. The CNN is trained with the Stochastic Gradient Descent (SGD) weights optimization method (we used a weight decay of 0.0001 and a momentum of 0.9).

For synthetic videos we calculate the NURD estimation error to evaluate the accuracy. For true videos, no guaranteed ground truth is available. We calculate the normalized Standard Deviation (STD) $\sigma$ for a certain time window of videos to access the stabilization performance [8]. The definition of STD is:

where ${\bar{\sigma}(f_{i,j})}$ is the STD calculated with pixel $f_{i,j}$ in one stacked frame stream, $i$ and $j$ are selected pixel indices on horizontal and vertical axis respectively. $N_{sig}$ is the number of pixels used to calculate $\bar{\sigma}$. Based on the STD calculation, we give a threshold to differentiate unstable pixels from stable pixels, and count the number of unstable pixels for each frame.

To evaluate the rotation error through a complete scan, we compute en-face projections [10, 11] of the OCT videos where each A-line is accumulated to one single value, so that the volumetric data in the polar domain are projected into 2D images. In this case the vertical Y axis corresponds to a circumferential scanning (B-Scan) and the horizontal X axis to a longitudinal volumetric scanning (3D Scan). By counting the shift of max intensity pixels in a whole stack, we can evaluate the precession angle. We also measure the local fluctuation on the en-face image , which represents the local shaking between neighboring frames.

The proposed stabilization algorithm is applied to both scanning modes. A pixel-wise stack smoothing is deployed to be the baseline method (average filtering of each pixel using 3 consecutive frames).

Table II shows the precession of the whole data stream, local fluctuation angle and mean pixel counts (MPC) of unstable pixels. For both scanning modes, the precession angle is reduced by 75%-80%, and for the robotic scan the precession is stabilised to only 6^∘. The baseline method can reduce the intensity derivation by blurring the 3D data, but does not improve the rotations and angular fluctuations. The proposed method not only reduces the overall rotations, but also reduces the pixel-wise instability without blurring the image or changing other information within every OCT frame (see the images in the subsection for qualitative results).

We use both the conventional internal pullback and the robotic scan and the synthetic video to study the functionality of the two parts of the algorithm (the element-wise NURD estimation part, and the overall rotation estimation part). We did a thorough ablation comparison by disabling any of those two parts, and the results are shown in Table III. The most detrimental modification is to disable the overall estimation part. In this case the proposed algorithm can no longer provide a valid angular value to compensate the drift error, and the precession error is no longer corrected. Sometimes it will even introduce a precession larger than the original precession, for example in the case of the real robotic scanning and synthetic videos. Disabling the element-wise NURD estimation also affects the reduction of the precession angle, but not as much as disabling the overall estimation. Compared to our full stabilization algorithm, disabling the NURD estimation obviously increases the unstable pixel count, because the overall rotation estimation can not align individual A-lines. This is worse for real scans where the overall rotation may have a larger estimation error due to the sparse features in sheath matching.

Fig. 6 shows qualitative results of pullback scans. To illustrate the NURD instability, we encode each 3 consecutive frames in independent channels of RGB color images[8]. In this RGB encoding, a colorful part of the image indicates the NURD artifact, and a well stabilized frame sequence should have minimal colorful pixels. We transform the data stream into the en-face projection (see description in section IV.A). All these formats of results are presented in Fig. 6. It can be seen that the “colorful” pixels of the original sequence (RGB encoded) is reduced by the proposed algorithm, and without changing the image quality. The 3D results are consistent with the en-face projections, where the angular rotation is reflected by the shift of high intensity pixels of en-face images. It can be seen that the stabilization algorithm straightens the high intensity lines of the original en-face projection, which maintains the intensity distribution of the whole data stream. More qualitative visual comparisons can be seen in Supplementary Movies.

Qualitative results on the ablation study are presented in fig. 7 and 8. Keeping only the NURD estimation will introduce iterative distortion to the image, and also limit the precession reduction (fig. 7 (B) and fig. 8 (B)). Fig. 7 (C) and fig. 8 (C) show the results with only the overall rotation estimation from the reference data stream. In this case the shaking of the frame sequence can be enlarged due to a large angular estimation error. The fusion stage from the full algorithm can effectively limit these phenomena.

In addition, we compare the NURD estimation part of the proposed CNN based method and a state-of-the-art approach based on the graph path searching (GS) [11], which is not a learning based method. Since this online correction method has no drift compensation, we disable the overall rotation estimation to compare the NURD estimation performance, using the indirect metrics mean pixel account (MPC) for in-stable pixels of real scans. We also compare them for the synthetic scans, using mean square error (MSE) of NURD estimation as a direct metric. Results from original scans and different algorithms are shown in fig. 9. Table IV shows statistics of these comparisons, the NURD estimation error of the proposed method is significantly lower than graph searching based method and can better reduce the number of unstable pixels (MPC).