Fast abdomen organ and tumor segmentation with nn-UNet

Yajun Wu, Ershuai Wang, Zhenzhou Shao

Keywords: Segmentation, nn-UNet, pseudo label

Abstract: The medical imaging community generates a wealth of datasets, many of which are openly accessible and annotated for specific diseases and tasks such as multi-organ or lesion segmentation. However, most datasets are only partially annotated for particular purpose, which hinders the training of multi-talent models. We uses a combination of pseudo labels and partial annotations to generate reliable fully annotated data, avoiding data conflict issues. Then, we designed a fast segmentation method for abdominal organs and tumors based on localization and segmentation. To accelerate inference, we adopt a slice-like downsample for location. To obtain the satisfactory segmentation, we first trained two models for organs and tumors with different target spacing, then combine the results. We also designed a weighted compound loss function and training patches selection strategy to finetuning the model. On the public validation set, the average scores of organ DSC, organ DSC, tumor DSC and tumor NSD are 0.9164, 0.9597, 0.4856 and 0.4221, respectively. Under our development enviroments, the average inference time is 8.54 seconds , the average maximum GPU memory is 4221.49 M, the average area under the GPU memory-time curve is 15074.59.

Submission Number: 4