Edge-Sensitive Left Ventricle Segmentation Using Deep Reinforcement Learning

Abstract

Deep reinforcement learning (DRL) has been utilized in numerous computer vision tasks, such as object detection, autonomous driving, etc. However, relatively few DRL methods have been proposed in the area of image segmentation, particularly in left ventricle segmentation. Reinforcement learning-based methods in earlier works often rely on learning proper thresholds to perform segmentation, and the segmentation results are inaccurate due to the sensitivity of the threshold. To tackle this problem, a novel DRL agent is designed to imitate the human process to perform LV segmentation. For this purpose, we formulate the segmentation problem as a Markov decision process and innovatively optimize it through DRL. The proposed DRL agent consists of two neural networks, i.e., First-P-Net and Next-P-Net. The First-P-Net locates the initial edge point, and the Next-P-Net locates the remaining edge points successively and ultimately obtains a closed segmentation result. The experimental results show that the proposed model has outperformed the previous reinforcement learning methods and achieved comparable performances compared with deep learning baselines on two widely used LV endocardium segmentation datasets, namely Automated Cardiac Diagnosis Challenge (ACDC) 2017 dataset, and Sunnybrook 2009 dataset. Moreover, the proposed model achieves higher F-measure accuracy compared with deep learning methods when training with a very limited number of samples.

Keywords: Markov decision process; deep reinforcement learning; double deep Q-network; image segmentation; left ventricle segmentation.

Figure 1

An example of the segmentation results of an image on Automated Cardiac Diagnosis Challenge (ACDC) 2017 training dataset across different epochs in the training process. The ground truth (GT) boundary is plotted in blue and the magenta dots are the points found by Next-P-Net. The red pentagram represents the first edge point found by First-P-Net.

Figure 2

The overall process of the proposed system: The First-P-Net finds the first edge point and generates a probability map of edge points positions. The Next-P-Net locates the next point based on the previous edge point and image information.

Figure 3

The left image defines the $nskipneighborhoods$ centered on the red point. The green points represent the eight skip neighborhoods of the red point. The middle image shows the defined action space and the corresponding action directions. The right image gives an example of the segmentation result. The ground truth (GT) boundary is plotted in blue and the magenta dots are the points found by Next-P-Net. The red pentagram represents the initial edge point.

Figure 4

Three separate reward functions: difference IoU reward function, edge distance reward function and points clustering reward function.

Figure 5

The architecture of First-P-Net. 3 × 3 conv or 1 × 1 conv: 3 × 3 or 1 × 1 convolution layer followed by batch normalization and ReLU activation function. resblk: revised ResNet Block. 2×: 2 upsampling. 4×: 4 upsampling. 0.5×: 0.5 downsampling. 1/8×: 1/8 downsampling.

Figure 6

The architecture of Next-P-Net. 7 × 7 conv: 7 × 7 convolution layer followed by batch normalization and ReLU activation function. resblk: ResNet Block.

Figure 7

Examples of segmentation outcomes. The first three rows are the segmentation performances on Sunnybrook 2009 testing dataset and the last three rows are the segmentation performances on ACDC 2017 testing dataset. The ground truth (GT) boundary is plotted in blue and the magenta dots are the points found by Next-P-Net. The red pentagram represents the first edge point found by First-P-Net.

Figure 8

Examples of the first edge point found by First-P-Net on ACDC 2017 testing dataset. The red pentagram represents the first edge point and the small image on the upper left or upper right corner is the partial enlargement of the first point.

Figure 9

States with different Q-values on ACDC 2017 testing Dataset. The first two rows show some images centered on the yellow point with low Q-value, while the last two rows show some images centered on the yellow point with high Q-value.

Figure 10

The changes in three separate reward values, total reward value, F-measure accuracy and APD accuracy according to the learning iterations in the training process on ACDC 2017 Dataset.

Figure 11

The changes in APD, F-measure, Precision and Recall according to the learning epochs in the training process on ACDC 2017 dataset.