|
""" |
|
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_hard_coral = 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)} corals detected:\n- {n_hard_coral} hard corals\n- {n_soft_coral} soft corals""" |
|
|
|
|
|
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) |
|
|
|
|
|
|
|
|
|
MODEL_FILEPATH_WEIGHTS = Path("data/model/best.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 benthic imagery segmentation 🪸", |
|
fn=fn, |
|
inputs=input, |
|
outputs=[output_image, output_raw], |
|
examples=image_filepaths, |
|
allow_flagging="never", |
|
) |
|
|
|
demo.launch() |
|
|
