MRL-Seg: Overcoming Imbalance in Medical Image Segmentation with Multi-Step Reinforcement Learning

Abstract

Medical image segmentation is a critical task for clinical diagnosis and research. However, dealing with highly imbalanced data remains a significant challenge in this domain, where the region of interest (ROI) may exhibit substantial variations across different slices. This presents a significant hurdle to medical image segmentation, as conventional segmentation methods may either overlook the minority class or overly emphasize the majority class, ultimately leading to a decrease in the overall generalization ability of the segmentation results. To overcome this, we propose a novel approach based on multi-step reinforcement learning, which integrates prior knowledge of medical images and pixel-wise segmentation difficulty into the reward function. Our method treats each pixel as an individual agent, utilizing diverse actions to evaluate its relevance for segmentation. To validate the effectiveness of our approach, we conduct experiments on four imbalanced medical datasets, and the results show that our approach surpasses other state-of-the-art methods in highly imbalanced scenarios. These findings hold substantial implications for clinical diagnosis and research.

@ARTICLE{10336383,
  author={Yang, Feiyang and Li, Xiongfei and Duan, Haoran and Xu, Feilong and Huang, Yawen and Zhang, Xiaoli and Long, Yang and Zheng, Yefeng},
  journal={IEEE Journal of Biomedical and Health Informatics}, 
  title={MRL-Seg: Overcoming Imbalance in Medical Image Segmentation With Multi-Step Reinforcement Learning}, 
  year={2024},
  volume={28},
  number={2},
  pages={858-869},
  keywords={Image segmentation;Transformers;Reinforcement learning;Lesions;Medical diagnostic imaging;Task analysis;Training;Deep learning;imbalanced medical image segmentation;radiomics;reinforcement learning},
  doi={10.1109/JBHI.2023.3336726}}