Segmentation of findings in the gastrointestinal tract is a challenging but also important task which is an important building stone for sufficient automatic decision support systems. In this work, we present our solution for the Medico 2020 task, which focused on the problem of colon polyp segmentation. We present our simple but efficient idea of using an augmentation method that uses grids in a pyramid-like manner (large to small) for segmentation. Our results show that the proposed methods work as indented and can also lead to comparable results when competing with other methods.
Our method has two main steps: data augmentation with PYRA using pre-defined grid sizes followed by training of a DL model with the resulting augmented data. The source code for our method can be found in our GitHub 1 repository. The development dataset [5]
Segmented polyp regions in Gastrointestinal Tract (GI) images [1] can provide detailed analysis to doctors to identify correct areas to proceed with treatments compared to other computer-aided analysis such as classification [2,9,10] and detection [7] which provide less detailed information about the exact region and size of the affected area. However, training Deep Learning (DL) models to perform segmentation for medical data is challenging because of the lack of medical domain images as a result of tight privacy restrictions, the high cost for annotating medical data using experts, and a lower number of true positive findings compared to true negatives. In this paper, we present our approach for the participation in the 2020 Medico Segmentation Challenge [4], for which we introduce a novel augmentation technique called pyramidfocus-augmentation (PYRA). PYRA can be used to improve the performance of segmentation tasks when we have a small dataset to train our DL models or if the number of positive findings is small. Further, our method can focus doctors' attention to regions of polyps gradually. In addition to that the output of the method is also adjustable meaning, we could present a lower resolution of the grid if this is sufficient for the task at hand which can help to save processing time. Finally, our technique can also be applied to any segmentation task using any deep learning segmentation model. provided by the organizers has 1000 polyp images with corresponding ground truth masks. We divided it into two parts such that 800 images are used for model training and 200 for testing.
As the first step in PYRA, we generate checker board grids as illustrated in the first row of Figure 2 with sizes of 𝑁 × 𝑁 with 𝑁 values of 4, 8, 16, 32, 64, 128 and 256. 𝑁 should be selected such that 𝑖𝑚𝑎𝑔𝑒_𝑠𝑖𝑧𝑒 % 𝑁 = 0. Applying these eight grid augmentations to the training dataset with 800 images increases the training data to 800 × 8 = 6400 images.
For the second step, we convert the Ground Truth (GT) segmentation masks into a grid-based representation of the GT corresponding to the grid sizes. For example, if the grid size is 8 × 8, then the corresponding GT is a 8 × 8 converted GT.
The transformation of the ground truth masks to gridded masks is performed as following: (i) we divide the gt into the input grid size, (ii) we counted true pixels of each grid cell, (iii) if the number of true pixels is larger than 0, we converted the whole cell into a true cell. An example of a converted GT is depicted on the top of Figure 1.
We have set up four experiments: Exp-1, Exp-2, Exp-3, and Exp-4 to show the performance of PYRA. Exp-1 and Exp-2 represent two baseline experiments. Exp-1 uses only the 800 training images without any augmentations. In Exp-2, we used general augmentations such as Affine, Coarse Dropout, and Additive Gaussian Noise from the library called imgaug [6]. Exp-3 and Exp-4 are using our PYRA with the data from Exp-1 and Exp-2, respectively. The training dataset size was changed from 800 to 6400 after applying PYRA. However, we validated our experiments only using 200 images reserved for testing. We have used one data loader for all experiments to maintain a fair evaluation. The baseline experiments Exp-1 and Exp-2 used the data loader with a grid size of 256 × 256 which represents the original GT masks without any conversion.
MediaEval'20, December 14-15 2020, Online We have used the Unet architecture [8] as our DL model to perform the polyp segmentation task. We trained the Unet model with a stacked input using a polyp image and a random grid image selected from the eight sizes. Then, the model was trained to predict converted GT which were formed by converting the real GT into a grid-based GT as in the previous section.
The Unet model used dropout layers with the probability of 0.5. Then, we used our Unet model as a stochastic model to perform Monte Carlo sampling for the validation data. We kept our Unet model in the training state to perform this sampling while predicting the output for the validation data. In the Pytorch library, which is used for all our implementations, we can do this simply by keeping the model state in the model.train() state. We iterated 50 times for a single input to predict the output. We calculated the mean from these 50 predictions, which is used as the final prediction for the competition and Standard Deviation (std) images to know the model's confidence for the predictions. The whole training process is illustrated in Figure 1 with an example image and a grid size of 8 × 8 as an input. However, we submitted the predicted mean images for the gird size of 256 × 256 which generate predictions with the size of true GT (without any transformations). All the experiments used a fixed learning rate of 0.001 with the RMSprop optimizer [3], which were selected from preliminary experiments.
Table 1 summarizes the Mean Intersection over Union (mIoU) and the Dice Coefficient (DC) for the validation dataset and the test dataset. The final results to the competition were collected from mean images calculated by sampling 50 outputs for the same input with the grid size of 256. Additionally, we have calculated std images for the validation dataset to show the benefits of using PYRA. Example outputs for a given input image are illustrated in Figure 2.
According to the results in Table 1, Exp-3 which use only Pyramidfocus-augmentation shows the best validation results with mIoU of 0.7693 and DC of 0.8447, and the best test results with mIoU of 0.6981 and DC of 0.7887. The advantage of our Pyramid-focusaugmentation can be identified using the third row of Figure 2 along the fourth row of the same figure. We can see that our model can focus on polyp regions step by step. The third row of Figure 2 shows how our model predicts correct polyp cells in 2 × 2, 4 × 4, 8 × 8, 16 × 16, 32 × 32, 64 × 64, 128 × 128 and 256 × 256 grid sizes, respectively. When we compare this row with the last row of the images of std, we can see that the model has high confidence for the identified polyp regions. For example, it shows high confidence (black color region) for the middle part of the polyps. In contrast, our model shows less confidence (yellow color region) for a polyps' outer borders.
In this paper, we presented a novel augmentation method called Pyramid-focus-augmentation (PYRA), which can be used to train segmentation DL methods. Our method shows a large benefit in the medical diagnosis use-case, by focusing a doctors' attention on regions with findings step by step.
Our experiments did not use post-processing to clean up output corresponding to the input grid. In future work, we will evaluate our approach with additional post-processing steps for smaller grid sizes. For example, we can do convolution operations to the output using a convolutional window equal to the input grid size to clean the results. However, post-processing techniques will not improve the final results when the grid size equals the input images' resolution.
The research has benefited from the Experimental Infrastructure for Exploration of Exascale Computing (eX3), which is financially supported by the Research Council of Norway under contract 270053.