Online detection of tomato defects based on YOLOv4 model pruning

Abstract: Surface defects have posed a negative impact on the quality and yield in the process of tomato growth. A post-harvest grading treatment can normally be utilized before tomato marketing. It is necessary to accurately and rapidly detect the defective tomato in the process of post-harvest grading. In this study, a real-time detection was proposed for the tomato surface defects using YOLOv4 model pruning. A diffuse light box was used to improve the acquisition system. High resolution images of tomatoes were then acquired to reduce the reflection of tomato surface under the direct exposure. Parallel computing and images combination were also selected for the high speed of image processing. The input images of the YOLOv4 network were generated to stitch the RGB images collected from three continuous detection stations. In addition, the channel pruning was selected to simplify the network parameters and structure in the YOLOv4 network model. There were a complex network structure and a large number of parameters in the original YOLOv4, leading to too large calculation and low inference speed of the model. The layer pruning was also used to further compress the depth of the network model on the basis of compressing the network width, in order to improve the detection speed for the real-time detection. A non-maximum suppression with the L1 norm was proposed to remove the redundant prediction box after fine-tuning network model, thereby accurately locating the defect location in the images. The detection performance of the improved model was evaluated using the YOLOv4 training data at the pruning rates under various target detection models. The maximum Mean Average Precision (mAP) value was taken as the pruning rate, indicating the minimum model size and inference time for the requirements of real-time performance. Therefore, the channel pruning rate of the YOLOv4 network was finally set to 80%, where the obtained model was named YOLOv4P. The YOLOv3, YOLOv4, YOLOv4-tiny, YOLOv4PC, and YOLOv4P models were compared to verify the performance of the detection model for the tomato defects. The results showed that the model pruning compression technology can effectively improve the detection speed at the lowest cost of accuracy. The real-time grading system was tested on the tomato experiment data set, including the stem, calyx, and defect types. The improved YOLOv4P network reduced the model size and reasoning time by 232.40 MB and 10.11 ms, respectively, compared with the original YOLOv4 network. In conclusion, the highest mAP and the fastest detection speed were achieved to detect the tomato surface defects using the improved model. The mAP increased from 92.45% to 94.56%, fully meeting the requirements of accurate and real-time detection. The finding can also provide an efficient online detection for the tomato real-time grading system.