RLSegNet: An Medical Image Segmentation Network Based on Reinforcement Learning.

Abstract

In the area of medical image segmentation, the spatial information can be further used to enhance the image segmentation performance. And the 3D convolution is mainly used to better utilize the spatial information. However, how to better utilize the spatial information in the 2D convolution is still a challenging task. In this paper, we propose an image segmentation network based on reinforcement learning (RLSegNet), which can translate the image segmentation process into a serial of decision-making problem. The proposed RLSegNet is a U-shaped network, which is composed of three components: the feature extraction network, the Mask Prediction Network (MPNet), and the up-sampling network with the cascade attention module. The deep semantic feature in the image is first extracted by adopting the feature extraction network. Then, the Mask Prediction Network (MPNet) is proposed to generate the prediction mask for the current frame based on the prior knowledge (segmentation result). And the proposed cascade attention module is mainly used to generate the weighted feature mask so that the up-sampling network pays more attention to the interesting region. Specifically, the state, action and reward used in the reinforcement learning are redesigned in the proposed RLSegNet to translate the segmentation process as the decision-making process, which performs as the reinforcement learning to realize the brain tumor segmentation. Extensive experiments are conducted on the BRATS 2015 dataset to evaluate the proposed RLSegNet. The experimental results demonstrate that the proposed method can achieve a better segmentation performance, in comparison with other state-of-the-art methods.