Spaces:
Running
Running
File size: 2,815 Bytes
c58ca54 c77fef7 c58ca54 c77fef7 c58ca54 c77fef7 c58ca54 c77fef7 c58ca54 c77fef7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
"""
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/")
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.launch()
|