feat: basic face recognition UI

app.py

@@ -10,29 +10,36 @@ import numpy as np
 from PIL import Image
 from ultralytics import YOLO
 
-from utils import (
-    bgr_to_rgb,
-    crop_from_yolov8,
-    get_best_device,
-    load_models,
-    resize,
-    run_pipeline,
-    setup,
-    square_pad,
-)
+from ui import bearid_ui
+from utils import load_models, run_pipeline, setup
 
 
-def prediction_to_str(yolo_prediction) -> str:
+def prediction_to_str(bear_ids: list[str], indexed_k_nearest_individuals) -> str:
     """
-    Turn the yolo_prediction into a human friendly string.
+    Turn the prediction into a human friendly string.
     """
-    boxes = yolo_prediction.boxes
-    return f"""{len(boxes.conf)} bear detected! Trigger the bear repellent 🐻"""
+    nearest_individuals = []
+
+    for j in range(len(bear_ids)):
+        bear_id = bear_ids[j]
+        nearest_individual = indexed_k_nearest_individuals[bear_id][0]
+        distance = nearest_individual["distance"]
+        nearest_individuals.append(
+            {"bear_id": bear_id, "distance": distance, "data": nearest_individual}
+        )
+
+    distance_str = "\n".join(
+        [
+            f"- {n['bear_id']} at distance {n['distance']:.2f} in the embedding space"
+            for n in nearest_individuals
+        ]
+    )
+    return f"The model found that the closest individual is {bear_ids[0]}:\n\n{distance_str}"
 
 
 def interface_fn(
     loaded_models: dict[str, Any], pil_image: Image.Image
-) -> Tuple[Image.Image, Image.Image, str]:
+) -> Tuple[Image.Image, Image.Image, Image.Image, str, list[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.
@@ -47,21 +54,80 @@ def interface_fn(
         model.
     """
     PARAM_SQUARE_DIM = 300
+    k = 5
+    n_samples_per_individuals = 4
     result = run_pipeline(
         loaded_models=loaded_models,
         pil_image=pil_image,
         param_square_dim=PARAM_SQUARE_DIM,
-        param_k=5,
-        param_n_samples_per_individual=4,
+        param_k=k,
+        param_n_samples_per_individual=n_samples_per_individuals,
         knn_index_filepath=METRIC_LEARNING_KNN_INDEX_FILEPATH,
     )
     pil_image_segmented_head = result["stages"]["segmentation"]["output"]["pil_image"]
     pil_image_cropped_head = result["stages"]["crop"]["output"]["pil_images"]["resized"]
 
-    # raw_prediction_str = prediction_to_str(yolov8_segmentation_prediction)
+    bear_ids = result["stages"]["identification"]["output"]["bear_ids"]
+    k_nearest_individuals = result["stages"]["identification"]["output"][
+        "k_nearest_individuals"
+    ]
+    indexed_samples = result["stages"]["identification"]["output"]["indexed_samples"]
+    indexed_k_nearest_individuals = result["stages"]["identification"]["output"][
+        "indexed_k_nearest_individuals"
+    ]
+    tmp = k_nearest_individuals
+    tmp = indexed_samples
+    bear_id = bear_ids[0]
+    other_bear_ids = bear_ids[1:]
+
+    images = []
+    for obid in other_bear_ids:
+        for fp in indexed_samples[obid]:
+            images.append(Image.open(fp))
+
+    output_filepath = Path("./data/07_model_output/predict/prediction.png")
+    output_filepath.parent.mkdir(parents=True, exist_ok=True)
+
+    bearid_ui(
+        pil_image_chip=pil_image_cropped_head,
+        indexed_k_nearest_individuals=indexed_k_nearest_individuals,
+        indexed_samples=indexed_samples,
+        save_filepath=output_filepath,
+        k_closest_neighbors=k,
+    )
+
+    pil_image_identification_prediction = Image.open(output_filepath)
 
-    return (pil_image_segmented_head, pil_image_cropped_head, str(result))
-    return (pil_image_segmented_head, raw_prediction_str)
+    nearest_individuals = []
+
+    for j in range(len(bear_ids)):
+        bear_id = bear_ids[j]
+        nearest_individual = indexed_k_nearest_individuals[bear_id][0]
+        distance = nearest_individual["distance"]
+        nearest_individuals.append(
+            {"bear_id": bear_id, "distance": distance, "data": nearest_individual}
+        )
+
+    distance_str = "\n".join(
+        [
+            f"- {n['bear_id']} at distance {n['distance']:.2f} in the embedding space"
+            for n in nearest_individuals
+        ]
+    )
+    tmp = f"The model found that the closest individual is {bear_id}:\n\n{distance_str}"
+
+    raw_prediction_str = prediction_to_str(
+        bear_ids=bear_ids, indexed_k_nearest_individuals=indexed_k_nearest_individuals
+    )
+
+    return (
+        pil_image_segmented_head,
+        pil_image_cropped_head,
+        pil_image_identification_prediction,
+        bear_id,
+        [Image.open(fp) for fp in indexed_samples[bear_id]],
+        raw_prediction_str,
+    )
 
 
 def examples(dir_examples: Path) -> list[Path]:
@@ -111,6 +177,13 @@ with gr.Blocks() as demo:
     )
     output_segmentation_image = gr.Image(type="pil", label="model prediction")
     output_cropped_image = gr.Image(type="pil", label="cropped bear face")
+    output_bear_id = gr.Text(label="Most likely BearID")
+    output_bear_id_samples = gr.Gallery(
+        label="Similar faces of the predicted BearID", preview=True
+    )
+    output_identification_prediction_image = gr.Image(
+        type="pil", label="prediction", scale=3
+    )
     output_raw = gr.Text(label="raw prediction")
 
     fn = lambda pil_image: interface_fn(
@@ -121,7 +194,14 @@ with gr.Blocks() as demo:
         title="ML pipeline for identifying bears from their faces 🐻",
         fn=fn,
         inputs=input,
-        outputs=[output_segmentation_image, output_cropped_image, output_raw],
+        outputs=[
+            output_segmentation_image,
+            output_cropped_image,
+            output_identification_prediction_image,
+            output_bear_id,
+            output_bear_id_samples,
+            output_raw,
+        ],
         examples=image_filepaths,
         flagging_mode="never",
     )

ui.py

@@ -0,0 +1,186 @@
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torchvision
+from matplotlib import font_manager
+from matplotlib.figure import Figure
+from matplotlib.gridspec import GridSpec
+from PIL import Image
+
+IMAGENET_MEAN = [0.485, 0.456, 0.406]
+IMAGENET_STD = [0.229, 0.224, 0.225]
+DISTANCE_THRESHOLD_NEW_INDIVIDUAL = 0.7
+
+
+def get_inverse_normalize_transform(mean, std):
+    return torchvision.transforms.Normalize(
+        mean=[-m / s for m, s in zip(mean, std)], std=[1 / s for s in std]
+    )
+
+
+def get_color(
+    distance: float,
+    distance_threshold_new_individual: float = DISTANCE_THRESHOLD_NEW_INDIVIDUAL,
+    margin: float = 0.10,
+) -> str:
+    threshold_unsure = distance_threshold_new_individual * (1.0 - margin)
+    threshold_new_individual = distance_threshold_new_individual * (1 + margin)
+    if distance < threshold_unsure:
+        return "green"
+    elif distance < threshold_new_individual:
+        return "orange"
+    else:
+        return "red"
+
+
+def draw_extrated_chip(ax, chip_image) -> None:
+    ax.set_title("Extracted chip")
+    ax.set_axis_off()
+    ax.imshow(chip_image)
+
+
+def draw_closest_neighbors(
+    fig: Figure,
+    gs: GridSpec,
+    i_start: int,
+    k_closest_neighbors: int,
+    indexed_k_nearest_individuals: dict,
+) -> None:
+    inv_normalize = get_inverse_normalize_transform(
+        mean=IMAGENET_MEAN,
+        std=IMAGENET_STD,
+    )
+
+    neighbors = []
+    for bear_id, xs in indexed_k_nearest_individuals.items():
+        for x in xs:
+            data = x.copy()
+            data["bear_id"] = bear_id
+            neighbors.append(data)
+
+    nearest_neighbors = sorted(
+        neighbors,
+        key=lambda x: x["distance"],
+    )[:k_closest_neighbors]
+    for j, neighbor in enumerate(nearest_neighbors):
+        ax = fig.add_subplot(gs[i_start, j])
+        distance = neighbor["distance"]
+        bear_id = neighbor["bear_id"]
+        dataset_image = neighbor["dataset_image"]
+        image = inv_normalize(dataset_image).numpy()
+        image = np.transpose(image, (1, 2, 0))
+        color = get_color(distance=distance)
+        ax.set_axis_off()
+        ax.set_title(label=f"{bear_id}: {distance:.2f}", color=color)
+        ax.imshow(image)
+
+
+def draw_top_k_individuals(
+    fig: Figure,
+    gs: GridSpec,
+    i_start: int,
+    i_end: int,
+    indexed_k_nearest_individuals: dict,
+    bear_ids: list[str],
+    indexed_samples: dict,
+):
+    inv_normalize = get_inverse_normalize_transform(
+        mean=IMAGENET_MEAN,
+        std=IMAGENET_STD,
+    )
+    for i in range(i_start, i_end):
+        for j in range(len(bear_ids)):
+            # Draw the closest individual chips
+            if i == i_start:
+                ax = fig.add_subplot(gs[i, j])
+                bear_id = bear_ids[j]
+                nearest_individual = indexed_k_nearest_individuals[bear_id][0]
+                distance = nearest_individual["distance"]
+                dataset_image = nearest_individual["dataset_image"]
+                image = inv_normalize(dataset_image).numpy()
+                image = np.transpose(image, (1, 2, 0))
+                color = get_color(distance=distance)
+                ax.set_axis_off()
+                ax.set_title(label=f"{bear_id}: {distance:.2f}", color=color)
+                ax.imshow(image)
+
+            # Draw random chips from the same individuals
+            else:
+                bear_id = bear_ids[j]
+                idx = i - i_start - 1
+                if idx < len(indexed_samples[bear_id]):
+                    filepath = indexed_samples[bear_id][idx]
+                    if filepath:
+                        ax = fig.add_subplot(gs[i, j])
+                        with Image.open(filepath) as image:
+                            ax.set_axis_off()
+                            ax.imshow(image)
+
+
+def bearid_ui(
+    pil_image_chip: Image.Image,
+    indexed_k_nearest_individuals: dict,
+    indexed_samples: dict,
+    save_filepath: Path,
+    k_closest_neighbors: int = 5,
+) -> None:
+    """Main UI for identifying bears."""
+    chip_image = pil_image_chip
+    # Assumption: the bear_ids are sorted by distance - if that's not something
+    # we can rely on, we should just sort
+    bear_ids = list(indexed_k_nearest_individuals.keys())
+
+    # Max of the number of closest_neighbors and the number of bearids
+    ncols = max(len(bear_ids), k_closest_neighbors)
+
+    # 1 row for the closest neighbors title section
+    # 1 row for the closest neighbors
+    # 1 row for the individuals title section
+    # rows for the indexed_samples (radom images of a given individual)
+    nrows = max([len(xs) for xs in indexed_samples.values()]) + 3
+    figsize = (3 * ncols, 3 * nrows)
+    fig = plt.figure(constrained_layout=True, figsize=figsize)
+    gs = GridSpec(nrows=nrows, ncols=ncols, figure=fig)
+    font_properties_section = font_manager.FontProperties(size=35)
+    font_properties_title = font_manager.FontProperties(size=40)
+
+    # Draw closest neighbors
+    i_closest_neighbors = 2
+    ax = fig.add_subplot(gs[i_closest_neighbors - 1, :])
+    ax.set_axis_off()
+    ax.text(
+        y=0.2,
+        x=0,
+        s="Closest faces",
+        font_properties=font_properties_section,
+    )
+    draw_closest_neighbors(
+        fig=fig,
+        gs=gs,
+        i_start=i_closest_neighbors,
+        k_closest_neighbors=k_closest_neighbors,
+        indexed_k_nearest_individuals=indexed_k_nearest_individuals,
+    )
+    # Filling out the grid with top k individuals and random samples
+    i_top_k_individual = 4
+    ax = fig.add_subplot(gs[i_top_k_individual - 1, :])
+    ax.set_axis_off()
+    ax.text(
+        y=0.2,
+        x=0,
+        s=f"Closest {len(bear_ids)} individuals",
+        font_properties=font_properties_section,
+    )
+    draw_top_k_individuals(
+        fig=fig,
+        gs=gs,
+        i_end=nrows,
+        i_start=i_top_k_individual,
+        indexed_k_nearest_individuals=indexed_k_nearest_individuals,
+        bear_ids=bear_ids,
+        indexed_samples=indexed_samples,
+    )
+
+    plt.savefig(save_filepath, bbox_inches="tight")
+    plt.close()