feat: initial setup, only the face segmentation stage
Browse files- .gitattributes +1 -0
- .gitignore +2 -0
- .python-version +1 -0
- __init__.py +0 -0
- app.py +111 -0
- data/09_external/artifacts/chips.zip +3 -0
- data/09_external/artifacts/packaged_pipeline.zip +3 -0
- data/images/P1250243.jpg +3 -0
- data/images/c2NhbGUoKQ.jpg +3 -0
- requirements.txt +4 -0
- scripts/chips/install.sh +14 -0
- utils.py +77 -0
.gitattributes
CHANGED
|
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
+
*.jpg filter=lfs diff=lfs merge=lfs -text
|
.gitignore
ADDED
|
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
|
|
|
| 1 |
+
data/06_models/
|
| 2 |
+
data/07_model_output/
|
.python-version
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
3.10.12
|
__init__.py
ADDED
|
File without changes
|
app.py
ADDED
|
@@ -0,0 +1,111 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Gradio app to showcase the pyronear model for early forest fire detection.
|
| 3 |
+
"""
|
| 4 |
+
|
| 5 |
+
from pathlib import Path
|
| 6 |
+
from typing import Tuple
|
| 7 |
+
|
| 8 |
+
import gradio as gr
|
| 9 |
+
import numpy as np
|
| 10 |
+
from PIL import Image
|
| 11 |
+
import subprocess
|
| 12 |
+
import shutil
|
| 13 |
+
import logging
|
| 14 |
+
import os
|
| 15 |
+
import torch
|
| 16 |
+
import pandas as pd
|
| 17 |
+
from ultralytics import YOLO
|
| 18 |
+
|
| 19 |
+
from utils import (
|
| 20 |
+
bgr_to_rgb,
|
| 21 |
+
get_best_device,
|
| 22 |
+
load_segmentation_model,
|
| 23 |
+
setup
|
| 24 |
+
)
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def prediction_to_str(yolo_prediction) -> str:
|
| 28 |
+
"""
|
| 29 |
+
Turn the yolo_prediction into a human friendly string.
|
| 30 |
+
"""
|
| 31 |
+
boxes = yolo_prediction.boxes
|
| 32 |
+
classes = boxes.cls.cpu().numpy().astype(np.int8)
|
| 33 |
+
n_bear = len([c for c in classes if c == 0])
|
| 34 |
+
n_soft_coral = len([c for c in classes if c == 1])
|
| 35 |
+
|
| 36 |
+
return f"""{len(boxes.conf)} bear detected! Trigger the bear repellent 🐻"""
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def interface_fn(model_segmentation: YOLO, pil_image: Image.Image) -> Tuple[Image.Image, str]:
|
| 40 |
+
"""
|
| 41 |
+
Main interface function that runs the model on the provided pil_image and
|
| 42 |
+
returns the exepected tuple to populate the gradio interface.
|
| 43 |
+
|
| 44 |
+
Args:
|
| 45 |
+
model_segmentation (YOLO): Loaded ultralytics YOLO model.
|
| 46 |
+
pil_image (PIL): image to run inference on.
|
| 47 |
+
|
| 48 |
+
Returns:
|
| 49 |
+
pil_image_with_prediction (PIL): image with prediction from the model.
|
| 50 |
+
raw_prediction_str (str): string representing the raw prediction from the
|
| 51 |
+
model.
|
| 52 |
+
"""
|
| 53 |
+
predictions = model_segmentation(pil_image)
|
| 54 |
+
prediction = predictions[0]
|
| 55 |
+
pil_image_with_prediction = Image.fromarray(bgr_to_rgb(prediction.plot()))
|
| 56 |
+
raw_prediction_str = prediction_to_str(prediction)
|
| 57 |
+
|
| 58 |
+
return (pil_image_with_prediction, raw_prediction_str)
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
def examples(dir_examples: Path) -> list[Path]:
|
| 62 |
+
"""
|
| 63 |
+
List the images from the dir_examples directory.
|
| 64 |
+
|
| 65 |
+
Returns:
|
| 66 |
+
filepaths (list[Path]): list of image filepaths.
|
| 67 |
+
"""
|
| 68 |
+
return list(dir_examples.glob("*.jpg"))
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
# Setting up the model artifacts
|
| 72 |
+
INPUT_PACKAGED_PIPELINE = Path("./data/09_external/artifacts/packaged_pipeline.zip")
|
| 73 |
+
PIPELINE_INSTALL_PATH = Path("./data/06_models/pipeline/metriclearning/")
|
| 74 |
+
|
| 75 |
+
setup(
|
| 76 |
+
input_packaged_pipeline=INPUT_PACKAGED_PIPELINE,
|
| 77 |
+
install_path=PIPELINE_INSTALL_PATH,
|
| 78 |
+
)
|
| 79 |
+
|
| 80 |
+
# Main Gradio interface
|
| 81 |
+
METRIC_LEARNING_MODEL_FILEPATH = Path("./data/06_models/pipeline/metriclearning/bearidentification/model.pt")
|
| 82 |
+
METRIC_LEARNING_KNN_INDEX_FILEPATH = Path("./data/06_models/pipeline/metriclearning/bearidentification/knn.index")
|
| 83 |
+
INSTANCE_SEGMENTATION_WEIGHTS_FILEPATH = Path("./data/06_models/pipeline/metriclearning/bearfacesegmentation/model.pt")
|
| 84 |
+
DIR_EXAMPLES = Path("data/images/")
|
| 85 |
+
DEFAULT_IMAGE_INDEX = 0
|
| 86 |
+
|
| 87 |
+
with gr.Blocks() as demo:
|
| 88 |
+
model_segmentation = load_segmentation_model(INSTANCE_SEGMENTATION_WEIGHTS_FILEPATH)
|
| 89 |
+
model_segmentation.info()
|
| 90 |
+
image_filepaths = examples(dir_examples=DIR_EXAMPLES)
|
| 91 |
+
default_value_input = Image.open(image_filepaths[DEFAULT_IMAGE_INDEX])
|
| 92 |
+
input = gr.Image(
|
| 93 |
+
value=default_value_input,
|
| 94 |
+
type="pil",
|
| 95 |
+
label="input image",
|
| 96 |
+
sources=["upload", "clipboard"],
|
| 97 |
+
)
|
| 98 |
+
output_image = gr.Image(type="pil", label="model prediction")
|
| 99 |
+
output_raw = gr.Text(label="raw prediction")
|
| 100 |
+
|
| 101 |
+
fn = lambda pil_image: interface_fn(model_segmentation=model_segmentation, pil_image=pil_image)
|
| 102 |
+
gr.Interface(
|
| 103 |
+
title="ML pipeline for identifying bears from their faces 🐻",
|
| 104 |
+
fn=fn,
|
| 105 |
+
inputs=input,
|
| 106 |
+
outputs=[output_image, output_raw],
|
| 107 |
+
examples=image_filepaths,
|
| 108 |
+
flagging_mode="never",
|
| 109 |
+
)
|
| 110 |
+
|
| 111 |
+
demo.launch()
|
data/09_external/artifacts/chips.zip
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:0a732292816cabd4e700abad2cbe856e56e921be3e91840dc62b926a50f4073f
|
| 3 |
+
size 144180439
|
data/09_external/artifacts/packaged_pipeline.zip
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:e471db5ac8c6abe07d3cf75f4b17382d3e35b2821075f4e6c1aa09bdcd9a7a62
|
| 3 |
+
size 284751540
|
data/images/P1250243.jpg
ADDED
|
Git LFS Details
|
data/images/c2NhbGUoKQ.jpg
ADDED
|
Git LFS Details
|
requirements.txt
ADDED
|
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
gradio==5.4.*
|
| 2 |
+
pandas==2.2.*
|
| 3 |
+
torch==2.5.*
|
| 4 |
+
ultralytics==8.3.*
|
scripts/chips/install.sh
ADDED
|
@@ -0,0 +1,14 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env bash
|
| 2 |
+
|
| 3 |
+
set -x
|
| 4 |
+
|
| 5 |
+
INPUT_CHIPS_PATH="./data/09_external/artifacts"
|
| 6 |
+
ZIP_ARCHIVE="${INPUT_CHIPS_PATH}/chips.zip"
|
| 7 |
+
OUTPUT_CHIPS_PATH="./data/07_model_output/bearfacesegmentation"
|
| 8 |
+
|
| 9 |
+
if [ ! -f "$ZIP_ARCHIVE" ]; then
|
| 10 |
+
echo "$ZIP_ARCHIVE does not exist."
|
| 11 |
+
exit 1
|
| 12 |
+
fi
|
| 13 |
+
|
| 14 |
+
unzip -f "$ZIP_ARCHIVE" -d "$OUTPUT_CHIPS_PATH"
|
utils.py
ADDED
|
@@ -0,0 +1,77 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from pathlib import Path
|
| 2 |
+
from typing import Tuple
|
| 3 |
+
|
| 4 |
+
import logging
|
| 5 |
+
import numpy as np
|
| 6 |
+
import os
|
| 7 |
+
import pandas as pd
|
| 8 |
+
import subprocess
|
| 9 |
+
import shutil
|
| 10 |
+
import torch
|
| 11 |
+
from ultralytics import YOLO
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
def get_best_device() -> torch.device:
|
| 15 |
+
"""Returns the best torch device depending on the hardware it is running
|
| 16 |
+
on."""
|
| 17 |
+
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
def _setup_chips() -> None:
|
| 21 |
+
"""
|
| 22 |
+
Setup the Database of chips used for the face recognition.
|
| 23 |
+
"""
|
| 24 |
+
subprocess.run(["./scripts/chips/install.sh"])
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def _setup_ml_pipeline(input_packaged_pipeline: Path, install_path: Path) -> None:
|
| 28 |
+
"""
|
| 29 |
+
Setup the ML pipeline, installing the model weights into their folders.
|
| 30 |
+
"""
|
| 31 |
+
logging.info(f"Installing the packaged pipeline in {install_path}")
|
| 32 |
+
os.makedirs(install_path, exist_ok=True)
|
| 33 |
+
packaged_pipeline_archive_filepath = input_packaged_pipeline
|
| 34 |
+
shutil.unpack_archive(
|
| 35 |
+
filename=packaged_pipeline_archive_filepath,
|
| 36 |
+
extract_dir=install_path,
|
| 37 |
+
)
|
| 38 |
+
metriclearning_model_filepath = install_path / "bearidentification" / "model.pt"
|
| 39 |
+
device = get_best_device()
|
| 40 |
+
bearidentification_model = torch.load(
|
| 41 |
+
metriclearning_model_filepath,
|
| 42 |
+
map_location=device,
|
| 43 |
+
)
|
| 44 |
+
df_split = pd.DataFrame(bearidentification_model["data_split"])
|
| 45 |
+
chips_root_dir = Path("/".join(df_split.iloc[0]["path"].split("/")[:-4]))
|
| 46 |
+
logging.info(f"Retrieved chips_root_dir: {chips_root_dir}")
|
| 47 |
+
os.makedirs(chips_root_dir, exist_ok=True)
|
| 48 |
+
shutil.copytree(
|
| 49 |
+
src=install_path / "chips",
|
| 50 |
+
dst=chips_root_dir,
|
| 51 |
+
dirs_exist_ok=True,
|
| 52 |
+
)
|
| 53 |
+
|
| 54 |
+
def setup(input_packaged_pipeline: Path, install_path: Path) -> None:
|
| 55 |
+
"""
|
| 56 |
+
Full setup of the project.
|
| 57 |
+
"""
|
| 58 |
+
_setup_chips()
|
| 59 |
+
_setup_ml_pipeline(
|
| 60 |
+
input_packaged_pipeline=input_packaged_pipeline,
|
| 61 |
+
install_path=install_path
|
| 62 |
+
)
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
def bgr_to_rgb(a: np.ndarray) -> np.ndarray:
|
| 66 |
+
"""
|
| 67 |
+
Turn a BGR numpy array into a RGB numpy array when the array `a` represents
|
| 68 |
+
an image.
|
| 69 |
+
"""
|
| 70 |
+
return a[:, :, ::-1]
|
| 71 |
+
|
| 72 |
+
|
| 73 |
+
def load_segmentation_model(filepath_weights: Path) -> YOLO:
|
| 74 |
+
"""
|
| 75 |
+
Load the YOLO model given the filepath_weights.
|
| 76 |
+
"""
|
| 77 |
+
return YOLO(filepath_weights)
|