Spaces:
Running
Running
| """ | |
| 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( | |
| fn=fn, | |
| inputs=input, | |
| outputs=[output_image, output_raw], | |
| examples=image_filepaths, | |
| flagging_mode="never", | |
| ) | |
| demo.launch() | |