feat: initial commit

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* 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
+*.jpg filter=lfs diff=lfs merge=lfs -text

.python-version

@@ -0,0 +1 @@
+3.10.12

app.py

@@ -0,0 +1,102 @@
+"""
+Gradio app to showcase the pyronear model for early forest fire detection.
+"""
+
+from pathlib import Path
+from typing import Tuple
+
+import gradio as gr
+import numpy as np
+from PIL import Image
+from ultralytics import YOLO
+
+
+def bgr_to_rgb(a: np.ndarray) -> np.ndarray:
+    """
+    Turn a BGR numpy array into a RGB numpy array when the array `a` represents
+    an image.
+    """
+    return a[:, :, ::-1]
+
+
+def prediction_to_str(yolo_prediction) -> str:
+    """
+    Turn the yolo_prediction into a human friendly string.
+    """
+    boxes = yolo_prediction.boxes
+    classes = boxes.cls.cpu().numpy().astype(np.int8)
+    n_bear = len([c for c in classes if c == 0])
+    n_soft_coral = len([c for c in classes if c == 1])
+
+    return f"""{len(boxes.conf)} bear detected! Trigger the bear repellent 🐻"""
+
+
+def predict(model: YOLO, pil_image: Image.Image) -> Tuple[Image.Image, str]:
+    """
+    Main interface function that runs the model on the provided pil_image and
+    returns the exepected tuple to populate the gradio interface.
+
+    Args:
+        model (YOLO): Loaded ultralytics YOLO model.
+        pil_image (PIL): image to run inference on.
+
+    Returns:
+        pil_image_with_prediction (PIL): image with prediction from the model.
+        raw_prediction_str (str): string representing the raw prediction from the
+        model.
+    """
+    predictions = model(pil_image)
+    prediction = predictions[0]
+    pil_image_with_prediction = Image.fromarray(bgr_to_rgb(prediction.plot()))
+    raw_prediction_str = prediction_to_str(prediction)
+
+    return (pil_image_with_prediction, raw_prediction_str)
+
+
+def examples(dir_examples: Path) -> list[Path]:
+    """
+    List the images from the dir_examples directory.
+
+    Returns:
+        filepaths (list[Path]): list of image filepaths.
+    """
+    return list(dir_examples.glob("*.jpg"))
+
+
+def load_model(filepath_weights: Path) -> YOLO:
+    """
+    Load the YOLO model given the filepath_weights.
+    """
+    return YOLO(filepath_weights)
+
+
+# Main Gradio interface
+
+MODEL_FILEPATH_WEIGHTS = Path("data/model/weights/model.pt")
+DIR_EXAMPLES = Path("data/images/")
+DEFAULT_IMAGE_INDEX = 1
+
+with gr.Blocks() as demo:
+    model = load_model(MODEL_FILEPATH_WEIGHTS)
+    image_filepaths = examples(dir_examples=DIR_EXAMPLES)
+    default_value_input = Image.open(image_filepaths[DEFAULT_IMAGE_INDEX])
+    input = gr.Image(
+        value=default_value_input,
+        type="pil",
+        label="input image",
+        sources=["upload", "clipboard"],
+    )
+    output_image = gr.Image(type="pil", label="model prediction")
+    output_raw = gr.Text(label="raw prediction")
+
+    fn = lambda pil_image: predict(model=model, pil_image=pil_image)
+    gr.Interface(
+        title="ML model for detecting bears from camera traps 🐻",
+        fn=fn,
+        inputs=input,
+        outputs=[output_image, output_raw],
+        examples=image_filepaths,
+        flagging_mode="never",
+    )
+
+demo.launch()

data/model/weights/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:47f41ed8d0867288d00d1ae823cb05ea90f412e81a66dd39939620f211a3eecd
+size 6346222

requirements.txt

@@ -0,0 +1,2 @@
+ultralytics==8.3.*
+gradio==5.4.*