Iteratively-Refined Interactive 3D Medical Image Segmentation with Multi-Agent Reinforcement Learning

TL;DR

This paper introduces a multi-agent reinforcement learning framework for iterative 3D medical image segmentation, modeling the process as an MDP and leveraging voxel-wise agents to improve accuracy and efficiency.

Contribution

It proposes a novel multi-agent RL approach that captures voxel dependencies and incorporates uncertainty, significantly enhancing interactive segmentation performance.

Findings

Outperforms state-of-the-art methods in accuracy.

Requires fewer user interactions for high-quality segmentation.

Achieves faster convergence in medical image segmentation tasks.

Abstract

Existing automatic 3D image segmentation methods usually fail to meet the clinic use. Many studies have explored an interactive strategy to improve the image segmentation performance by iteratively incorporating user hints. However, the dynamic process for successive interactions is largely ignored. We here propose to model the dynamic process of iterative interactive image segmentation as a Markov decision process (MDP) and solve it with reinforcement learning (RL). Unfortunately, it is intractable to use single-agent RL for voxel-wise prediction due to the large exploration space. To reduce the exploration space to a tractable size, we treat each voxel as an agent with a shared voxel-level behavior strategy so that it can be solved with multi-agent reinforcement learning. An additional advantage of this multi-agent model is to capture the dependency among voxels for segmentation task. Meanwhile, to enrich the information of previous segmentations, we reserve the prediction uncertainty in the state space of MDP and derive an adjustment action space leading to a more precise and finer segmentation. In addition, to improve the efficiency of exploration, we design a relative cross-entropy gain-based reward to update the policy in a constrained direction. Experimental results on various medical datasets have shown that our method significantly outperforms existing state-of-the-art methods, with the advantage of fewer interactions and a faster convergence.