Differentiable Projection from Optical Coherence Tomography B-Scan without Retinal Layer Segmentation Supervision

We conducted experiments on the public OCTA-500 [12] dataset, specifically the OCTA_6M subset, which contains 300 subjects with 6mm $\times$ 6mm FOV. As demonstrated in Fig. 2, each sample is paired OCT/OCTA B-scans and projection maps, where B2/B3 are projected from OCT and B5/B6 are from OCTA. Each slice of B-scans corresponds to a column in the projection maps. B2/B3 are the focus in this study. B2, projected by averaging pixels between ILM and OPL layers (i.e., the red and green curves in Fig. 4), shows the vessels in the inner retina with high reflection; and B3, averagely projected between OPL and BM layers (i.e., the green and blue curves in Fig. 4), shows the vessel shadows in the outer retina with low reflection. We use the official split, with 180/20/100 subjects (72,000/8,000/40,000 B-scans) for training, validation and evaluation, respectively.

As depicted in Fig. 1, we first use a 2D U-Net [14] at a B-scan level to output coordinates of retinal layers, which are then processed by the proposed Differentiable Projection Module (DPM) to generate the projection maps. As the ground truth of the target PMs is min-max normalized over the whole PMs, we design a Conditional Min-Max (CMM) normalization trick during training to rescale the intensity values. The model parameters are trained via loss back-propagation between the prediction and ground truth of PMs.

Given an B-scan of $H\times W$, we use a downsampled image of $H_{2}\times W_{2}$ ($H_{2}=H/2,W_{2}=W/2$) as network input to reduce memory. The 2D U-Net backbone produces a feature map of $H_{2}\times W_{2}$. We apply a $H_{2}\times 1$ convolution with 3 output channels followed by a horizontal bilinear upsampling layer. The $3\times W$ output, representing predicted coordinates of the 3 retinal layers (i.e., ILM, OPL and BM) are mapped into $[-1,1]$ with $\mathit{tanh}$.

For B2/B3, the projection maps are generated by averaging pixels between 2 layers (Sec. 2.1). For simplicity, we describe the Differentiable Projection Module (DPM) for 2 layers (i.e., $2\times W$). As shown in Fig. 3, we uniformly sample $M$ points between the 2 layers (including the end points), which represent the vertical coordinates of sampled areas ($M\times W$). Then we assign a abscissa coordinate to each vertical coordinate, composing a spatial position matrix $G\in\mathbb{R}^{M\times W\times 2}$. For each coordinate, we could bilinearly interpolate the pixel value from the B-scan. This pixel interpolation operation introduced by the spatial transformer networks [13] could be easily implemented in modern deep learning frameworks, e.g., $\mathit{grid\_sample}$ in PyTorch [15]. The sampled output of $M\times W$ represents the warped image of pixels between the 2 layers. It could be averaged pooled vertically into $1\times W$, which is a column in the target projection map.

Our model produces a target PM column-by-column (i.e., slice-by-slice in B-scans); However, the PM ground truth provided in the OCTA-500 Dataset [12] is min-max normalized over the whole PMs. The accurate min-max values of prediction could be only obtained given all slices, which is yet impossible during training due to memory constraint. Inspired by the auto-decoding (NOT auto-encoding) training technique [16], we propose a Conditional Min-Max normalization trick to rescale the intensity of the projection maps:

$$\mathit{CMM}_{i}(I)=\frac{I-I_{i}^{min}}{I_{i}^{max}-I_{i}^{min}},$$

(1)

where $I_{i}^{max}$ and $I_{i}^{min}$ are learnable parameters for the maximum and minimum values of the projected B-scans, conditional on the training subject ID $i$. Specifically, there are $180\times 2$ parameters for 180 training subjects in our study.

The loss function of our model is composed of 2 terms: an L-1 loss $L_{1}$ and a feature loss $L_{F}$. The $L_{1}$ measures the pixel-wise similarity, while the $L_{F}$ measures structural similarity by computing the feature similarity using a pretrained network [17]. The project loss for B2 is defined as

$$L_{B_{2}}=L_{1}(y_{i},\mathit{CMM}_{i}(f(x_{i})))+L_{F}(y_{i},\mathit{CMM}_{i}(f(x_{i}))),$$

(2)

where $y_{i}$ is the B2 ground truth, $f(x_{i})$ is the pre-normalzied prediction. When training the B₂ and B₃ projection maps simultaneously, the final loss is

$$L=L_{B_{2}}+\lambda L_{B_{3}}.$$

(3)

The $\mathit{CMM}$ normalization trick is only used during training. During inference, we predict pre-normalized outputs slice-by-slice, and then normalize it using the real min-max values on the predicted PMs.

To validate the effectiveness of the proposed methods, we design the baselines: 1) CNN only: a vanilla CNN without DPM. 2) DPM only: an optimization-based DPM without the CNN part as a deep prior, i.e., optimizing the layer coordinates directly on each data pair. We use PSNR and SSIM to assess the predicted projection maps, where PSNR focuses on pixel similarity and SSIM focues on structural similarity.

Our framework is implemented in PyTorch [15], all networks were trained on 4 NVIDIA RTX 2080 Ti with an Adam optimizer [18] using a batch size of 72. The learning rate starts with $1\times 10^{-4}$ and exponentially decays with a ratio of $1\times 10^{-2}$ after every epoch. We use $\lambda=0.2$ in Eq. 3 to balance the training for B₂ and B₃ with a simple grid search.

We quantitatively compared the performance of our method against CNN-only method and DPM-only method to assess the effect of each component of our model, we use PSNR and SSIM for evaluation. Table 1 gives a quantitative comparison of the results. From Table 1, we can observe that DPM-only method has the poorest performance, for this model has no component to understanding the input B-scan globally, just adjusts the layer segmentations blindly on the basis of the target projection map. CNN-only method also has a bad performance, we assume that this is because this method just uses a CNN network to predict the projection directly, ignoring the correct layered structure to get the target projection map.

In addition, we assessed the performance of CMM. The comparison demonstrates that adding CMM module is beneficial for learning layered structural information. However, the improvement of SSIM in B3 is not obvious, we assume that the reason is that B3 displays the vessel shadows in the outer retina with low reflection, in which vessel information is the dominance and the edges of tissues are distinct, so the accuracy of predicting ILM and BM layers does not have a great impact on projection B3 to a certain extent.

In Fig. 4, we show projection maps obtained from CNN-only, DPM-only, our method (CNN+DPM) and without-CMM, and layers predicted by our method. The CNN-only model can produce most blood vessels successfully but has serious distortion in details. The DPM-only model totally fails to produce a projection map, the image is full of noise for the predicted layers have crossed, but the contours of some thick blood vessels are visible. From the last three images, we can clearly see that adding CMM module can help produce a clearer projection map and get more precise layers.

In Fig. 5 (a), we illustrate two cases with lesions, which are difficult to predict layers. In these cases, the boundaries of layers are blurred due to the presence of the diseased areas, which bring difficulties to the segmentation task. Our model tries to learn structural feature of diseased areas, and get acceptable results.

In Fig. 5 (b), we illustrate two typical failure cases. In the first example, our model fails to correctly segment the bottom layer of the B-scan, result in a serious distortion of the bottom part of the predicted projection map. The reason is that the bottom right part of the poor-quality original image lacks structural information. In the second example, the layers predicted by our model can not fit the target layers in B-scan very well, lead to generating a poor projection. The feature of the original B-scan is that the imaging of the retina is steep, which is rare in the used dataset, leading our model to learn more about flat B-scans and trying to predict smoother layers.

In this section, we try to transfer the segmentations to the OCTA B-scans. OCTA is non-invasive technique that shows details of blood vessels that have low inherent contrast in OCT images. Given that OCTA images are generated from the OCT images, we directly transfer the layers of OCT to the OCTA, no need to train the model on OCTA again, and generate good quality projection maps. In Fig. 6, we illustrate two cases to show the results of transformation. Fig. 6 (b) shows the projection results of OCTA, we can see that the synthesized projection map generates the shape of blood vessels greatly and is structurally coherent with its corresponding ground-truth.