In this paper, we present our method for Medico automatic polyp segmentation challenge at MediaEval 2020. In our method, we utilized the knowledge distillation technique to improve ResUNet++ which performs well on automatic polyp segmentation. In our experiment, our proposed model called KD-ResUNet++ outperforms ResUNet++ in terms of Jaccard index, Dice similarity coeficient, and recall. Our best models achieved Jaccard index, Dice similarity coeficient, and FPS of 0.6196, 0.7089, and 107.8797 respectively on the oficial test dataset in the challenge.
Automatic polyp segmentation is a challenging task due to
variations in the shape and size of polyps. In this paper, we propose
KD-ResUNet++, which is based on the ResUNet++ architecture [
U-Net is a very popular deep learning architecture for biomedical
image segmentation. U-Net won the 2015 ISBI cell tracking
challenge. ResUNet [
Knowledge distillation aims at transferring dark knowledge from a
teacher model such as wide [
In this section, the architecture of our proposed method for automatic polyp segmentation is first presented. After that, we describe the details of the key components in the following subsections. The overall architecture of our proposed model is shown in Figure 1. First, input images are augmented by the data augmentation module. Second, augmented images are used as the input of both the student model and the teacher model. Third, distillation loss between the output of the student model and the output of the teacher model and the student loss between the output of the student model and ground-truth label is calculated to train the student model. 3.1
Deep learning models require a large amount of training data to work efectively. However, the size of the provided colonoscopy dataset is not very large. To solve this problem, data augmentation can be used to make the relatively smaller dataset a large one. It is reported that the performance of the deep learning model can be improved by augmenting the existing data rather than collecting new data. In our data augmentation step, we used 2 augmentation strategies (rotation and horizontal flipping) to generate new training sets. The rotation operation used for data augmentation is done by randomly rotating the input by 90 degrees zero or more times. The rotation operation fills the area of rotated images where there was no image pixel with black. In addition, we applied horizontal lfipping to each of the rotated images. In our proposed approach, we use self-knowledge distillation where teacher network and student network are the same. We use ResUNet++ for both teacher and student networks. In knowledge distillation, the teacher network is first trained to transfer knowledge to the student network. Also, the loss function consists of 1) the distillation loss and 2) the student loss. The distillation loss can be calculated using dice loss between the output of the student model and the output of pre-trained teacher model , and the student loss can be calculated using dice loss between the output of the student model and the ground-truth label as follows: = 1 − ( , ) = 1 − ( , ) = 0.1 ∗ + The total loss is then calculated as the joint of the distillation and student losses as follows: (1) (2) (3) 4 4.1
We trained our proposed model using the Kvasir-SEG dataset [
The dataset is split into 88 % for learning the weights and 12 % for validating the model during the training step. Before the training step, we augment input images using our data augmentation module described in Section 3.1 and converted the images to the size of 256×256 pixels. The validation set is only normalized. The learning rate is set to 0.001. We use Adam as our optimizer. Our results on the validation set are presented in Table 1. Also, the oficial results on polyp segmentation and algorithm eficiency tasks on the same test dataset are shown in Table 2 and Table 3, respectively. Table 1 shows that ResUNet++ achieved slightly better precision than KD-ResUNet++. However, KD-ResUNet++ outperforms ResUNet++ in terms of Jaccard index and Dice similarity coeficient which is an important metric for semantic segmentation task. Table 1 shows that our proposed model outperforms ResUNet++ in terms of Jaccard index, Dice similarity coeficient, and recall. Table 2 and 3 show that our proposed model achieved Jaccard index, Dice similarity coeficient, and FPS of 0.6196, 0.7089, and 107.8797 respectively on the oficial test dataset. Besides, examples of three diferent segmentations produced by ResUNet++ and KD-ResUNet++ are depicted in Figure 2. Figure 2 shows that the result of KD-ResUNet++ are more similar with ground truth than the result of ResUNet++. 5
In this paper, we presented KD-ResUNet++ for automatic polyp
segmentation. In our proposed framework, the data augmentation
technique is applied to input images. Also, we use self-knowledge
distillation where teacher and student networks are the same. We
use the ResUNet++ model for our student and teacher networks.
Our proposed model is evaluated on the validation set as well as the
oficial test set. Our experimental results show that our proposed
model outperforms ResUNet++ in terms of Jaccard index, Dice
similarity coeficient, and recall. These results indicate that our
proposed method can capture polyp segmentation boundary well
and could be potentially used in clinical settings. In the future, we
plan to use diferent knowledge types in our knowledge distillation.
Also, we plan to modify the ResUNet++ architecture to incorporate
the model pre-trained on a large image dataset (e.g., ImageNet [
This research was supported by the Basic Science Research Program through the National Research Foundation of Korea (NRF) funded by the Ministry of Education (Grant No. NRF-2020R1I1A3074141).