Online Hard Patch Mining Using Shape Models and Bandit Algorithm for Multi-Organ Segmentation

Hard sample selection can effectively improve model convergence by extracting the most representative samples from a training set. However, due to the large capacity of medical images, existing sampling strategies suffer from insufficient exploitation for hard samples or high time cost for sample selection when adopted by 3D patch-based models in the field of multi-organ segmentation. In this paper, we present a novel and effective online hard patch mining (OHPM) algorithm. In our method, an average shape model that can be mapped with all training images is constructed to guide the exploration of hard patches and aggregate feedback from predicted patches. The process of hard mining is formalized as a multi-armed bandit problem and solved with bandit algorithms. With the shape model, OHPM requires negligible time consumption and can intuitively locate difficult anatomical areas during training. The employment of bandit algorithms ensures online and sufficient hard mining. We integrate OHPM with advanced segmentation networks and evaluate them on two datasets containing different anatomical structures. Comparative experiments with other sampling strategies demonstrate the superiority of OHPM in boosting segmentation performance and improving model convergence. The results in each dataset with each network suggest that OHPM significantly outperforms other sampling strategies by nearly 2% average Dice score.