improving documentation

README.md

@@ -34,29 +34,40 @@ base_model: Ultralytics/YOLOv8
 ---
 [Croppie](https://croppie.org/) cherry detection model © 2024 by [Alliance Bioversity & CIAT](https://alliancebioversityciat.org/), [Producers Direct](https://producersdirect.org/) and [M-Omulimisa](https://m-omulimisa.com/) is licensed under [GNU-GPLv3](https://www.gnu.org/licenses/quick-guide-gplv3.html)
 
+**Funded by**: Deutsche Gesellschaft für Internationale Zusammenarbeit (GIZ) [Fair Forward Initiative - AI for All](https://huggingface.co/fair-forward)
+
 ## General description
-Ultralytics' Yolo V8 medium model fined tuned for coffee cherry detection using the [Croppie coffee dataset](https://huggingface.co/datasets/rgautroncgiar/croppie_coffee_ug).
+Ultralytics' Yolo V8 medium [model fined tuned](https://yolov8.org/how-to-use-fine-tune-yolov8/) for coffee cherry detection using the [Croppie coffee dataset](https://huggingface.co/datasets/rgautroncgiar/croppie_coffee_ug). 
 
-![](images/annotated_1688033955437_.jpg)
-
-**Note: the low visibility/unsure class was not used for model fine tuning**
+This algorithm provides automated cherry count from RGB pictures. Takes as input a picture and returns the cherry count by class.
 
 The predicted numerical classes correspond to the following cherry types:
 ```
 {0: "dark_brown_cherry", 1: "green_cherry", 2: "red_cherry", 3: "yellow_cherry"}
 ```
 
+**Examples of use**:
+* yield estimates
+* ripeness detection
+
+**Limitations:** This algorithm does not include correction of cherry occlusion.
+
+![](images/annotated_1688033955437_.jpg)
+
+**Note: the low visibility/unsure class was not used for model fine tuning**
+
+
 ## Repository structure
 
 ```
 .
 ├── images
-│   ├── 1688033955437.jpg  # image for test
-│   └── annotated_1688033955437_.jpg
+│   ├── foo.bar # images for the documentation
 ├── model_v3_202402021.pt  # fine tuning of Yolo v8
 ├── README.md
 ├── LICENSE.txt # detailed term of the software license
 └── scripts
+    ├── custom_training.py # script which overwrites the default YOLO training method
     ├── render_results.py  # helper function to annotate predictions
     ├── requirements.txt  # pip requirements
     └── test_script.py  # test script
@@ -67,6 +78,20 @@ Assuming you are in the ```scripts``` folder, you can run ```python3 test_script
 
 Make sure that the Python packages found in ```requirements.txt``` are installed. In case they are not, simply run ```pip3 install -r requirements.txt```.
 
+A live demonstration is freely accesible [here](https://croppie.org/).
+
+## Training metrics
+![](images/training_results.png)
+
+The model has been trained using the custom YOLO class found in ```./scripts/custom_YOLO.py```. The custom YOLO class can be exactly used as the original [YOLO class](https://docs.ultralytics.com/reference/models/yolo/model/). The hyperparameters used during the training can be found in ```./scripts/args.yaml```.
+
+## Test metrics
+
+<img src="images/F1_curve.png" width="300">
+<img src="images/P_curve.png" width="300">
+<img src="images/PR_curve.png" width="300">
+<img src="images/R_curve.png" width="300">
+
 ## License
 [Croppie](https://croppie.org/) cherry detection model © 2024 by [Alliance Bioversity & CIAT](https://alliancebioversityciat.org/), [Producers Direct](https://producersdirect.org/) and [M-Omulimisa](https://m-omulimisa.com/) is licensed under [GNU-GPLv3](https://www.gnu.org/licenses/quick-guide-gplv3.html)
 
@@ -77,3 +102,7 @@ This program is distributed in the hope that it will be useful, but WITHOUT ANY
 You should have received a copy of the GNU General Public License along with this program. If not, see <https://www.gnu.org/licenses/>.
 
 The detailed terms of the license are available in the ```LICENSE``` file in the repository.
+
+## Funding
+
+**Funded by**: Deutsche Gesellschaft für Internationale Zusammenarbeit (GIZ) [Fair Forward Initiative - AI for All](https://huggingface.co/fair-forward)

scripts/test_script.py

@@ -21,7 +21,10 @@ if __name__=='__main__':
                     'red_cherry': '#F44336',
                     'dark_brown_cherry': '#C36105',
                     'low_visibility_unsure': '#02D5FA'}
-    class_dic = {1: 'green_cherry', 3: 'yellow_cherry', 2: 'red_cherry', 0: 'dark_brown_cherry'}
+    class_dic = {1: 'green_cherry',
+                 3: 'yellow_cherry',
+                 2: 'red_cherry',
+                 0: 'dark_brown_cherry'}
     render = annotate_image_prediction(
         image_path=image_path,
         yolo_boxes=results[0].boxes,