Sequential Modelling for Medical Image Segmentation

Abstract

An image can be seen as a long sequence of tokens with spatial structure hence an image segmentation task can be treated as a sequence-to-sequence task. Existing attention-based segmentation models have incorporated an up-sampling module at the pixel decoder, such design assumed the backbone to have multiple feature scales. Furthermore, the existing models did not emphasize the ability of sequential modelling for multi-tasking. In this paper, we propose to model image segmentation task as a long sequence prediction task with an improvisation on MegaByte. As MegaByte efficiently reduced the self-attention cost using multi-scale structure known as global and local models, we train two MegaByte models as an encoder and decoder for multi-tasking. Our multi-task model allows different segmentation tasks to be trained simultaneously and inferenced with a task prompt. The encoder takes in the image sequence to encode them as a key while the decoder takes in the conditional input as a query. As decoder conditional input is usually sparser than the image, we use a relatively shallow decoder. We demonstrate our method on segmentation of intracranial haemorrhages from computed tomography (CT) head scans and show comparable performance to existing models. On a brain tissue segmentation task on CT images, our method outperformed the existing methods especially on tiny structure. Clinical relevanceCT images are the first line of head scans when stroke or haemorrhage is suspected. Segmentation of intracranial haemorrhages (ICH) can provide clinicians with important measures of brain lesions to decide on treatment procedure or surgical decision. Brain and tissue volumes are indicative of neurodegenerative diseases. The proposed techniques can be used to segment ICH and brain tissues from CT head scans.