Add confidence-based keypoint outputs

Extend the app and UI to surface keypoint confidence more clearly. This adds optional confidence-based coloring with a legend, an annotated image download, and a confidence table sorted by lowest scores first. It also preloads the default SuperAnimal PyTorch model in a background thread so the first request is less likely to block on model download.

app.py

@@ -4,6 +4,7 @@
 # Adapted from  https://huggingface.co/spaces/DeepLabCut/MegaDetector_DeepLabCut
 
 import os
+import threading
 import yaml
 import numpy as np
 from matplotlib import cm
@@ -16,9 +17,10 @@ import dlclive
 from PIL import Image, ImageColor, ImageFont, ImageDraw
 
 from viz_utils import save_results_as_json, draw_keypoints_on_image, draw_bbox_w_text, save_results_only_dlc, save_results_pytorch
+from viz_utils import add_confidence_legend, keypoint_confidence_rows, save_annotated_image
 from detection_utils import predict_md, crop_animal_detections
 from dlc_utils import predict_dlc
-from pytorch_utils import predict_superanimal, PYTORCH_MODELS
+from pytorch_utils import predict_superanimal, load_superanimal, PYTORCH_MODELS
 from ui_utils import gradio_inputs_for_MD_DLC, gradio_outputs_for_MD_DLC, gradio_description_and_examples
 
 from deeplabcut.utils import auxiliaryfunctions
@@ -46,6 +48,16 @@ DLC_models_dict = {'superanimal_topviewmouse': ('superanimal_topviewmouse_dlcrne
 
 
 
+#####################################################
+def finalize_outputs(img_output, download_file, kpts_per_animal, map_label_id_to_str, color_by_confidence):
+    # confidence legend, annotated image for download and per-keypoint confidence table
+    if color_by_confidence:
+        img_output = add_confidence_legend(img_output)
+    annotated_file = save_annotated_image(img_output)
+    confidence_rows = keypoint_confidence_rows(kpts_per_animal, map_label_id_to_str)
+    return img_output, download_file, annotated_file, confidence_rows
+
+
 #####################################################
 def predict_pipeline_pytorch(img_input,
                              superanimal,
@@ -57,6 +69,7 @@ def predict_pipeline_pytorch(img_input,
                              font_size,
                              keypt_color,
                              marker_size,
+                             flag_color_by_confidence,
                              ):
     # detection + pose with the SuperAnimal PyTorch models (keypoints in image coords)
     img_output, animals, bodyparts = predict_superanimal(img_input,
@@ -75,7 +88,8 @@ def predict_pipeline_pytorch(img_input,
                                 font_style=font_style,
                                 font_size=font_size,
                                 keypt_color=keypt_color,
-                                marker_size=marker_size)
+                                marker_size=marker_size,
+                                color_by_confidence=flag_color_by_confidence)
         if not flag_dlc_only:
             draw_bbox_w_text(img_output,
                              animal['bbox'],
@@ -84,7 +98,9 @@ def predict_pipeline_pytorch(img_input,
     pose_model, detector = PYTORCH_MODELS[superanimal]
     download_file = save_results_pytorch(animals, map_label_id_to_str, superanimal,
                                          pose_model, None if flag_dlc_only else detector)
-    return img_output, download_file
+    return finalize_outputs(img_output, download_file,
+                            [animal['kpts'] for animal in animals], map_label_id_to_str,
+                            flag_color_by_confidence)
 
 
 #####################################################
@@ -100,6 +116,7 @@ def predict_pipeline(img_input,
                      font_size,
                      keypt_color,
                      marker_size,
+                     flag_color_by_confidence,
                      ):
 
     if backend == "PyTorch":
@@ -112,7 +129,8 @@ def predict_pipeline(img_input,
                                         font_style,
                                         font_size,
                                         keypt_color,
-                                        marker_size)
+                                        marker_size,
+                                        flag_color_by_confidence)
 
     # TensorFlow (legacy): MegaDetector crops + DLCLive
     dlc_model_name, dlc_model_dir = DLC_models_dict[dlc_model_input_str]
@@ -169,11 +187,14 @@ def predict_pipeline(img_input,
                                 font_style=font_style,
                                 font_size=font_size,
                                 keypt_color=keypt_color,
-                                marker_size=marker_size)
+                                marker_size=marker_size,
+                                color_by_confidence=flag_color_by_confidence)
 
         donw_file = save_results_only_dlc(list_kpts_per_crop[0], map_label_id_to_str,dlc_model_name)
 
-        return img_input, donw_file
+        return finalize_outputs(img_input, donw_file,
+                                [list_kpts_per_crop[0]], map_label_id_to_str,
+                                flag_color_by_confidence)
 
     else:
         # Compute kpts for each crop
@@ -202,7 +223,8 @@ def predict_pipeline(img_input,
                                     font_style=font_style,
                                     font_size=font_size,
                                     keypt_color=keypt_color,
-                                    marker_size=marker_size)
+                                    marker_size=marker_size,
+                                    color_by_confidence=flag_color_by_confidence)
 
             # Paste crop in original image
             img_background.paste(img_crop,
@@ -217,7 +239,9 @@ def predict_pipeline(img_input,
         # Save detection results as json
         download_file  = save_results_as_json(md_results,list_kpts_per_crop,list_bboxes,map_label_id_to_str,dlc_model_name,mega_model_input)
 
-        return img_background, download_file
+        return finalize_outputs(img_background, download_file,
+                                list_kpts_per_crop, map_label_id_to_str,
+                                flag_color_by_confidence)
 
 
 
@@ -254,5 +278,9 @@ with gr.Blocks(title=gr_title) as demo:
                 cache_examples=True,
                 cache_mode="lazy")
 
+# download and build the default model while the app starts; a request arriving
+# earlier waits on the same lock instead of downloading again
+threading.Thread(target=load_superanimal, args=("superanimal_quadruped",), daemon=True).start()
+
 demo.queue(default_concurrency_limit=1)  # PyTorch runners are not thread-safe
 demo.launch(theme=gr.themes.Default())

ui_utils.py

@@ -57,6 +57,11 @@ def gradio_inputs_for_MD_DLC(backends_list, md_models_list, dlc_models_list):
             label="Show bodypart labels?",
         )
 
+        gr_color_by_confidence_checkbox = gr.Checkbox(
+            value=True,
+            label="Color keypoints by confidence? (otherwise by bodypart)",
+        )
+
         gr_keypt_color = gr.ColorPicker(
             value="#862db7",
             label="Choose color for keypoint label",
@@ -98,28 +103,34 @@ def gradio_inputs_for_MD_DLC(backends_list, md_models_list, dlc_models_list):
         gr_slider_font_size,
         gr_keypt_color,
         gr_slider_marker_size,
+        gr_color_by_confidence_checkbox,
     ]
 
 
 def gradio_outputs_for_MD_DLC():
     gr_image_output = gr.Image(type="pil", label="Output Image")
-    gr_file_download = gr.File(label="Download JSON file")
-    return [gr_image_output, gr_file_download]
+    with gr.Row():
+        gr_file_download = gr.File(label="Download JSON file")
+        gr_image_download = gr.File(label="Download annotated image")
+    gr_confidence_table = gr.Dataframe(
+        headers=["animal", "bodypart", "confidence"],
+        label="Keypoint confidence (lowest first)",
+        interactive=False,
+    )
+    return [gr_image_output, gr_file_download, gr_image_download, gr_confidence_table]
 
 
 def gradio_description_and_examples():
-    title = "DeepLabCut Model Zoo SuperAnimals"
+    title = "DeepLabCut Model Zoo: SuperAnimals"
     description = (
-        "Test the SuperAnimal models from the "
-        "<a href='http://www.mackenziemathislab.org/dlc-modelzoo'>"
-        "DeepLabCut ModelZoo Project</a>, and read more on arXiv: "
-        "https://arxiv.org/abs/2203.07436! Simply upload an image and see how it does. "
-        "Want to run on videos on the cloud or locally? See the "
-        "<a href='http://www.mackenziemathislab.org/dlc-modelzoo'>DeepLabCut ModelZoo</a>."
+        "Estimate animal poses with the SuperAnimal models from the "
+        "[DeepLabCut Model Zoo](http://www.mackenziemathislab.org/dlc-modelzoo) "
+        "([paper](https://arxiv.org/abs/2203.07436)). "
+        "Upload an image or pick an example below; to run on videos, see the Model Zoo page."
     )
 
     examples = [
-        [image, "PyTorch", "md_v5a", "superanimal_quadruped", False, True, 0.5, 0.4, "amiko", 10, "#ff0000", 5]
+        [image, "PyTorch", "md_v5a", "superanimal_quadruped", False, True, 0.5, 0.4, "amiko", 10, "#ff0000", 5, True]
         for image in ("examples/dog.jpeg", "examples/cat.jpg")
     ]
 

viz_utils.py

@@ -25,6 +25,7 @@ def draw_keypoints_on_image(image,
                             font_size=8,
                             keypt_color="#ff0000",
                             marker_size=2,
+                            color_by_confidence=True,
                             ):
     """Draws keypoints on an image.
     Modified from:
@@ -47,14 +48,7 @@ def draw_keypoints_on_image(image,
     im_width, im_height = image.size
     keypoints_x = [k[0] for k in keypoints]
     keypoints_y = [k[1] for k in keypoints]
-    alpha = [k[2] for k in keypoints]
-    norm = matplotlib.colors.Normalize(vmin=0, vmax=255)
-
-    # debugging keypoints
-    print (keypoints)
-                                
-    names_for_color = [i for i in map_label_id_to_str.keys()]
-    colores = np.linspace(0, 255, num=len(names_for_color),dtype= int)
+    confidences = [k[2] for k in keypoints]
 
     # adjust keypoints coords if required
     if use_normalized_coordinates:
@@ -64,15 +58,19 @@ def draw_keypoints_on_image(image,
     #cmap = matplotlib.cm.get_cmap('hsv')
     # draw ellipses around keypoints
     for i, (keypoint_x, keypoint_y) in enumerate(zip(keypoints_x, keypoints_y)):
-        round_fill = list(cm.viridis(norm(colores[i]),bytes=True))#[round(num*255) for num in list(cmap(i))[:3]] #check!
         # handling potential nans in the keypoints
         if np.isnan(keypoint_x).any():
             continue
-        
-        if np.isnan(alpha[i]) == False : 
-            round_fill[3] = round(alpha[i] *255)
-        #print(round_fill)
-        #round_outline = [round(num*255) for num in list(cmap2(alpha[i]))[:3]]
+
+        confidence = float(np.clip(confidences[i], 0, 1))
+        if color_by_confidence:
+            # fill color encodes the keypoint confidence (see add_confidence_legend)
+            round_fill = cm.viridis(confidence, bytes=True)
+        else:
+            # one color per bodypart, transparency encodes the confidence
+            round_fill = list(cm.viridis(i / max(len(keypoints) - 1, 1), bytes=True))
+            round_fill[3] = round(confidence * 255)
+            round_fill = tuple(round_fill)
         draw.ellipse([(keypoint_x - marker_size, keypoint_y - marker_size),
                       (keypoint_x + marker_size, keypoint_y + marker_size)],
                        fill=tuple(round_fill), outline= 'black', width=1) #fill and outline: [0,255]
@@ -86,6 +84,55 @@ def draw_keypoints_on_image(image,
                       ImageColor.getcolor(keypt_color, "RGB"), # rgb #
                       font=font)
 
+#########################################
+#  Legend for the keypoint confidence colors
+def add_confidence_legend(image, font_style='amiko'):
+    """Returns the image with a white band below it holding a viridis strip (confidence 0 to 1).
+
+    The band is added below the image, so the legend never covers it and keypoint
+    coordinates are unchanged.
+    """
+    im_width, im_height = image.size
+    strip_w = max(60, im_width // 5)
+    strip_h = max(6, im_height // 60)
+    font = ImageFont.truetype(FONTS[font_style], max(10, strip_h * 2))
+    margin = strip_h
+    band_h = strip_h + font.size + 3 * margin
+
+    out = Image.new("RGB", (im_width, im_height + band_h), "white")
+    out.paste(image.convert("RGB"), (0, 0))
+    draw = ImageDraw.Draw(out)
+    x0 = im_width - strip_w - 2 * margin
+    y0 = im_height + margin
+    for dx in range(strip_w):
+        draw.line([(x0 + dx, y0), (x0 + dx, y0 + strip_h)],
+                  fill=cm.viridis(dx / (strip_w - 1), bytes=True))
+    label_y = y0 + strip_h + margin // 2
+    draw.text((x0, label_y), "0", fill="black", font=font)
+    draw.text((x0 + strip_w, label_y), "1", fill="black", font=font, anchor="ra")
+    draw.text((x0 + strip_w / 2, label_y), "confidence", fill="black", font=font, anchor="ma")
+    return out
+
+
+#########################################
+#  Keypoint confidences as table rows
+def keypoint_confidence_rows(kpts_per_animal, map_label_id_to_str):
+    """(animal, bodypart, confidence) for every keypoint kept (not NaN), lowest confidence first."""
+    rows = []
+    for i_animal, kpts in enumerate(kpts_per_animal):
+        for i_kpt, kpt in enumerate(kpts):
+            if not np.isnan(kpt[2]):
+                rows.append([i_animal, map_label_id_to_str[i_kpt], round(float(kpt[2]), 3)])
+    return sorted(rows, key=lambda row: row[2])
+
+
+#########################################
+#  Save the annotated image for download
+def save_annotated_image(image, path_to_output_file='download_annotated.png'):
+    image.save(path_to_output_file)
+    return path_to_output_file
+
+
 #########################################
 #  Draw bboxes on image
 def draw_bbox_w_text(img,