Enhancing Polyp Segmentation Generalizability by Minimizing Images' Total Variation

Polyps are considered a precursor of colon cancer and early detection of polyps may help decrease mortality rate. Several deep learning models have been proposed to address the problem, however, with limited generalizability due to the scarcity of the current public datasets. To tackle the issue, researchers typically use data augmentation techniques or generative models to inflate training images, independent of a downstream learning task. In this paper, we propose a deep learning framework to jointly train an image transformation model with a segmentation model where the output of the former is the input of the latter. During training, the image transformation model generates variations of the input image at every epoch, implicitly increasing the training data size for the segmentation model. On the other hand, we design a total variational denoising cost for the image transformation model, which effectively ensures that a transformation applied to an input image works towards the segmentation and not any other random effects which may hurt the segmentation goal. The experimental results with different settings demonstrate that the proposed framework has consistently shown an improvement of approximately 1% to 10% polyp IoU on unseen test images.