feat: implement interface for forest fire detector ML model

app.py

@@ -1,9 +1,97 @@
+"""
+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
+    return f"""{len(boxes.conf)} fire smoke(s) detected!\nwith bounding boxes (xyxy): {boxes.xyxy.cpu().numpy()}\nand confidence scores: {boxes.conf.cpu().numpy()}"""
+
+
+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/best.pt")
+DIR_EXAMPLES = Path("data/images/")
 
-def greet(name):
-    return "Hello " + name + "!!"
+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[1])
+    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 early forest fire detection 🔥",
+        fn=fn,
+        inputs=input,
+        outputs=[output_image, output_raw],
+        examples=image_filepaths,
+        allow_flagging="never",
+    )
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
 demo.launch()