Polyp-YOLOv5-Tiny: A Lightweight Model for Real-Time Polyp Detection
Shimin Ou, Yixing Gao, Zebin Zhang, Chenjian Shi
Abstract
Early colonoscopy diagnosis can significantly reduce the mortality rate of colon cancer patients. Deep learning based object detection assists to enhance clinical performance on accurate diagnosis; however, it is challenging to integrate current object detection models on low-performance endoscope hardware devices. This paper presents a lightweight model for real-time polyp detection. By reducing the number of convolutional kernels by half and removing the large-object-detecting head from YOLOv5, our model got a similar precision to YOLOv3-spp (mAP.5:.95 of 0.591 and 0.583 respectively), but with a model size of only 2.8 MB (119.7 MB of YOLOv3-spp). There is a slight loss in precision compared to YOLOv5s (mAP.5:.95 of 0.636 with 13.7MB model size); however, our model still shows a significant advantage on reducing model complexity. Experimental results also indicated that MS COCO pre-training is helpful in polyp detection tasks, and the Mosaic augmentation strongly enhances the precision of YOLO models especially when the training set is small.