Initial commit

.gitattributes

@@ -0,0 +1,39 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+images/annotated_1688033955437_.jpg filter=lfs diff=lfs merge=lfs -text
+images/1688033955437_.jpg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+**/_*
+_*

README.md

@@ -0,0 +1,52 @@
+---
+tags:
+- coffee
+- cherry count
+- yield estimate
+- ultralyticsplus
+- yolov8
+- ultralytics
+- yolo
+- vision
+- object-detection
+- pytorch
+
+library_name: ultralytics
+library_version: 8.0.75
+inference: false
+
+datasets:
+- rgautron/croppie_coffee
+
+model-index:
+- name: rgautron/croppie_coffee
+  results:
+  - task:
+      type: object-detection
+
+    dataset:
+      type: rgautron/croppie_coffee
+      name: croppie_coffee
+      split: val
+
+    metrics:
+      - type: precision  # substitute for [email protected]
+        value: 0.691
+        name: [email protected](box)
+---
+### 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_split).
+
+![](images/annotated_1688033955437_.jpg)
+
+**Note: the low visibility/unsure class was not used for model fine tuning**
+
+The predicted numerical classes correspond to the following cherry types:
+```
+{0: "dark_brown_cherry", 1: "green_cherry", 2: "red_cherry", 3: "yellow_cherry"}
+```
+
+### Demonstration
+Assuming you are in the ```scripts``` folder, you can run ```python3 test_script.py```. This script saves the annotated image in ```../images/annotated_1688033955437.jpg```.
+
+Make sure that the Python packages found in ```requirements.txt``` are installed. In case they are not, simply run ```pip3 install -r requirements.txt```.

model_v3_202402021.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3180ebfc4422978e38d18722ba23840debaf531867770aaaddebba2f6431b988
+size 52013835

scripts/render_results.py

@@ -0,0 +1,71 @@
+import cv2
+import matplotlib.pyplot as plt
+from PIL import ImageColor
+from pathlib import Path
+import os
+
+def annotate_image_prediction(image_path, yolo_boxes, class_dic, saving_folder, hex_class_colors=None, show=False, true_count=False, saving_image_name=None, put_title=True, box_thickness=3, font_scale=1, font_thickness=5):
+    """
+    Fonction to label individual images with YOLO predictions
+    Args:
+        image_path (str): path to the image to label
+        yolo_boxes (str): YOLO predicted boxes
+        class_dic (dict): dictionary with predicted class as key and corresponding label as value
+        saving_folder (str): folder where to save the annotated image
+        hex_class_colors (dict, optional): HEX color code dict of the class to plot. Defaults to None.
+        show (bool, optional): If you want a window of the annotated image to pop up. Defaults to False.
+        true_count (bool, optional): If you want to display the true total count of cherries. Defaults to None.
+        saving_image_name (str, optional): Name of the annotated image to save. Defaults to None.
+        put_title (bool, optional): If you want a title to show in the plot. Defaults to True.
+        box_thickness (int, optional): Thickness of the bounding boxes to plot. Defaults to 3.
+        font_scale (int, optional): Font scale of the text of counts to be displayed. Defaults to 1.
+        font_thickness (int, optional): Font thickness of the text of counts to be displayed. Defaults to 5.
+
+    Returns:
+        string: saving path of the annotated image
+    """    
+    if os.path.isfile(image_path):
+        Path(saving_folder).mkdir(parents=True, exist_ok=True)
+        image_file = image_path.split('/')[-1]
+        if not hex_class_colors:
+            hex_class_colors = {class_name: (255, 0, 0) for class_name in class_dic.values()}
+        color_map = {key: ImageColor.getcolor(hex_class_colors[class_dic[key]], 'RGB') for key in [*class_dic]}
+
+        img = cv2.imread(image_path)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        dh, dw, _ = img.shape
+
+        for yolo_box in yolo_boxes:
+            x1, y1, x2, y2 = yolo_box.xyxy[0]
+            x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
+            c = int(yolo_box.cls[0])
+            cv2.rectangle(img, (x1, y1), (x2, y2), color_map[c], box_thickness)
+
+        if show:
+            plt.imshow(img)
+            plt.show()
+        img_copy = img.copy()
+        if put_title:
+            if true_count:
+                title = f'Predicted count: {len(yolo_boxes)}, true count: {true_count}, delta: {len(yolo_boxes) - true_count}'
+            else:
+                title = f'Predicted count: {len(yolo_boxes)}'
+            cv2.putText(
+                img=img_copy,
+                text=title,
+                org=(int(0.1 * dw), int(0.1 * dh)),
+                fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+                fontScale=font_scale,
+                thickness=font_thickness,
+                color=(255,251,5),
+            )
+
+        if not saving_image_name:
+            saving_image_name = f'annotated_{image_file}'
+        Path(saving_folder).mkdir(parents=True, exist_ok=True)
+        full_saving_path = os.path.join(saving_folder, saving_image_name)
+        plt.imsave(full_saving_path, img_copy)
+    else:
+        full_saving_path = None
+        print(f'WARNING: {image_path} does not exists')
+    return full_saving_path

scripts/requirements.txt

@@ -0,0 +1,5 @@
+torch==2.1.2
+ultralytics==8.0.75
+Pillow==9.3.0
+matplotlib==3.7.1
+opencv-python==4.7.0.72

scripts/test_script.py

@@ -0,0 +1,34 @@
+import ultralytics
+ultralytics.checks()
+from render_results import annotate_image_prediction
+from PIL import Image
+
+if __name__=='__main__':
+    # load model
+    model = ultralytics.YOLO('../model_v3_202402021.pt')
+    
+    # set parameters
+    image_path = '../images/1688033955437.jpg'
+    saving_folder = '../images'
+    
+    # infer
+    image = Image.open(image_path)
+    results = model.predict(image_path)
+    
+    # plot
+    hex_class_colors = {'green_cherry': '#9CF09A',
+                    'yellow_cherry': '#F3C63D',
+                    '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'}
+    render = annotate_image_prediction(
+        image_path=image_path,
+        yolo_boxes=results[0].boxes,
+        class_dic=class_dic,
+        saving_folder=saving_folder,
+        hex_class_colors=hex_class_colors,
+        show=True,
+        font_scale=5,
+        font_thickness=10,
+        )