feat: initial setup, only the face segmentation stage

.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

.gitignore

@@ -0,0 +1,2 @@
+data/06_models/
+data/07_model_output/

.python-version

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

app.py

@@ -0,0 +1,111 @@
+"""
+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
+import subprocess
+import shutil
+import logging
+import os
+import torch
+import pandas as pd
+from ultralytics import YOLO
+
+from utils import (
+    bgr_to_rgb,
+    get_best_device,
+    load_segmentation_model,
+    setup
+)
+
+
+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 interface_fn(model_segmentation: 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_segmentation (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_segmentation(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"))
+
+
+# Setting up the model artifacts
+INPUT_PACKAGED_PIPELINE = Path("./data/09_external/artifacts/packaged_pipeline.zip")
+PIPELINE_INSTALL_PATH = Path("./data/06_models/pipeline/metriclearning/")
+
+setup(
+    input_packaged_pipeline=INPUT_PACKAGED_PIPELINE,
+    install_path=PIPELINE_INSTALL_PATH,
+)
+
+# Main Gradio interface
+METRIC_LEARNING_MODEL_FILEPATH = Path("./data/06_models/pipeline/metriclearning/bearidentification/model.pt")
+METRIC_LEARNING_KNN_INDEX_FILEPATH = Path("./data/06_models/pipeline/metriclearning/bearidentification/knn.index")
+INSTANCE_SEGMENTATION_WEIGHTS_FILEPATH = Path("./data/06_models/pipeline/metriclearning/bearfacesegmentation/model.pt")
+DIR_EXAMPLES = Path("data/images/")
+DEFAULT_IMAGE_INDEX = 0
+
+with gr.Blocks() as demo:
+    model_segmentation = load_segmentation_model(INSTANCE_SEGMENTATION_WEIGHTS_FILEPATH)
+    model_segmentation.info()
+    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: interface_fn(model_segmentation=model_segmentation, pil_image=pil_image)
+    gr.Interface(
+        title="ML pipeline for identifying bears from their faces 🐻",
+        fn=fn,
+        inputs=input,
+        outputs=[output_image, output_raw],
+        examples=image_filepaths,
+        flagging_mode="never",
+    )
+
+demo.launch()

data/09_external/artifacts/chips.zip

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

data/09_external/artifacts/packaged_pipeline.zip

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

requirements.txt

@@ -0,0 +1,4 @@
+gradio==5.4.*
+pandas==2.2.*
+torch==2.5.*
+ultralytics==8.3.*

scripts/chips/install.sh

@@ -0,0 +1,14 @@
+#!/usr/bin/env bash
+
+set -x
+
+INPUT_CHIPS_PATH="./data/09_external/artifacts"
+ZIP_ARCHIVE="${INPUT_CHIPS_PATH}/chips.zip"
+OUTPUT_CHIPS_PATH="./data/07_model_output/bearfacesegmentation"
+
+if [ ! -f "$ZIP_ARCHIVE" ]; then
+	echo "$ZIP_ARCHIVE does not exist."
+	exit 1
+fi
+
+unzip -f "$ZIP_ARCHIVE" -d "$OUTPUT_CHIPS_PATH"

utils.py

@@ -0,0 +1,77 @@
+from pathlib import Path
+from typing import Tuple
+
+import logging
+import numpy as np
+import os
+import pandas as pd
+import subprocess
+import shutil
+import torch
+from ultralytics import YOLO
+
+
+def get_best_device() -> torch.device:
+    """Returns the best torch device depending on the hardware it is running
+    on."""
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def _setup_chips() -> None:
+    """
+    Setup the Database of chips used for the face recognition.
+    """
+    subprocess.run(["./scripts/chips/install.sh"])
+
+
+def _setup_ml_pipeline(input_packaged_pipeline: Path, install_path: Path) -> None:
+    """
+    Setup the ML pipeline, installing the model weights into their folders.
+    """
+    logging.info(f"Installing the packaged pipeline in {install_path}")
+    os.makedirs(install_path, exist_ok=True)
+    packaged_pipeline_archive_filepath = input_packaged_pipeline
+    shutil.unpack_archive(
+        filename=packaged_pipeline_archive_filepath,
+        extract_dir=install_path,
+    )
+    metriclearning_model_filepath = install_path / "bearidentification" / "model.pt"
+    device = get_best_device()
+    bearidentification_model = torch.load(
+        metriclearning_model_filepath,
+        map_location=device,
+    )
+    df_split = pd.DataFrame(bearidentification_model["data_split"])
+    chips_root_dir = Path("/".join(df_split.iloc[0]["path"].split("/")[:-4]))
+    logging.info(f"Retrieved chips_root_dir: {chips_root_dir}")
+    os.makedirs(chips_root_dir, exist_ok=True)
+    shutil.copytree(
+        src=install_path / "chips",
+        dst=chips_root_dir,
+        dirs_exist_ok=True,
+    )
+
+def setup(input_packaged_pipeline: Path, install_path: Path) -> None:
+    """
+    Full setup of the project.
+    """
+    _setup_chips()
+    _setup_ml_pipeline(
+        input_packaged_pipeline=input_packaged_pipeline,
+        install_path=install_path
+    )
+
+
+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 load_segmentation_model(filepath_weights: Path) -> YOLO:
+    """
+    Load the YOLO model given the filepath_weights.
+    """
+    return YOLO(filepath_weights)