PR-YOLO: Improved YOLO for fast protozoa classification and segmentation

Abstract

Protozoa, such as Ceratium and Paramecium, play a fundamental role in establishing sustainable ecosystems. The distribution and classification of certain protozoa and their species are informative indicators to evaluate environmental quality. However, protozoa analysis is traditionally performed by molecular biological (DNA, RNA) or morphological methods, which are time-consuming and require an experienced laboratory operator. In this work, we adopt a deep learning-based network to solve the protozoa classification task. This method utilizes microscope images to help researchers analyse the protozoa population and species, reducing the cost of experimental sample storage and relieving the burden on laboratory operators. However, the shape and size of protozoa vary greatly, which places a great burden on the optimization of DCNN feature distillation. It is a great challenge to build a fast and precise protozoa analysis image. We present a new version of YOLO-v5 with better performance and extend it with instance segmentation called PR-YOLO. Building on the original YOLOv5, we added two extra parallel branches to PR-YOLO, which perform different segmentation subtasks: (1) a branch generates a set of prototype masks (images); (2) the other branch predicts a set of mask coefficients corresponding to prototype masks for each instance mask generation. Then, to improve the classification accuracy, we introduced transformer encoding blocks and lightweight Convolution Block Attention Modules (CBAMs) to explore the prediction potential with a self-attention mechanism. To quantitatively evaluate the performance of PR-YOLO, a comprehensive experiment was carried out on the hand-segmented microscopic protozoa images. Our model obtained the best results, with average classification accuracy of 96.83% and mean Average Precision(mAP) of 86.92% with a speed of 25.2 fps, which proves that the method has high robustness in this application field.