For the paper 'Real-time mango detection on CPU using pruned YOLO network'
A Keras implementation of pruned YOLOv3-tiny to detect mango data (Tensorflow backend). Original YOLO implement is inspired by qqwweee/keras-yolo3.
-
Processing your data. The mango dataset is opensource can be found in here. Mango dataset labels are all xml files. So voc_annotation.py is for you to transfer them into txt file for our code.
-
The COCO2017 dataset can be found here. COCO website offer a lot of API for quickly using the COCO data. You can easily get the apple and orange images as you need.
-
Download YOLOv3-tiny pre-trained weigths from YOLO website. Then convert.py can be used to transfer the .weights file to .h5 file.
-
Now the GradAM (gradient of target output with respect to the activation maps * activation maps) can be computed using compute_grad_am.py. We also upload the results in model_data/grad_am_sort_idx_L.npy, so you can directly use it to prune the YOLOv3-tiny.
-
Some description for code in main dir.
- kmeans.py is for computing anchor sizes.
- compute_grad_am.py is used for computing F1 scores of all network.
- compute_flops.py is used for computing FLOPs in Table 1 of the manuscript.
-
Then we will use the GradAM to prune the network. Using "prune_for_retrain_yolo.py" to prune the original YOLO network, then using "retrain_pruned_net.py" to finetuning it.
-
"train_MangoYolo.py" is for reproducing paper of MangoYOLO. The trained weights and network structure can be find in logs for validate the results in our manuscript.
-
Some description for directories.
- data_annotation stores formatted txt file of mango dataset.
- dataset stores COCO apple and orange images and mango dataset.
- logs stores trained network weights.
- yolo3 contains all network building codes and some utils.
- For the Table 1 in the manuscript
- After reformatting the mango data labels to .txt, run detect_mango_cfg.py to generate the F1 score in Table 1. You can change param cfg to use other trained networks to detect the mango.
@article{shi2020attribution,
title={An attribution-based pruning method for real-time mango detection with YOLO network},
author={Shi, Rui and Li, Tianxing and Yamaguchi, Yasushi},
journal={Computers and Electronics in Agriculture},
volume={169},
pages={105214},
year={2020},
publisher={Elsevier}
}