Polish pose visualization UI

Refresh the Gradio app with a DeepLabCut-themed UI, configurable keypoint colormaps and bounding-box colors, and a standalone HTML confidence legend instead of drawing it into the output image. This also improves example defaults for different image sizes, cleans up displayed bodypart names, and makes bbox labels more readable with contrast-aware text.

app.py

@@ -20,9 +20,14 @@ from PIL import Image
 from detection_utils import crop_animal_detections, predict_md
 from dlc_utils import predict_dlc
 from pytorch_utils import PYTORCH_MODELS, load_superanimal, predict_superanimal
-from ui_utils import gradio_description_and_examples, gradio_inputs_for_MD_DLC, gradio_outputs_for_MD_DLC
+from ui_utils import (
+    confidence_legend_html,
+    dlc_theme,
+    gradio_description_and_examples,
+    gradio_inputs_for_MD_DLC,
+    gradio_outputs_for_MD_DLC,
+)
 from viz_utils import (
-    add_confidence_legend,
     draw_bbox_w_text,
     draw_keypoints_on_image,
     keypoint_confidence_rows,
@@ -53,13 +58,11 @@ DLC_models_dict = {
 
 
 #####################################################
-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)
+def finalize_outputs(img_output, download_file, kpts_per_animal, map_label_id_to_str, color_by_confidence, colormap):
     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
+    legend = confidence_legend_html(colormap) if color_by_confidence else ""
+    return img_output, legend, download_file, annotated_file, confidence_rows
 
 
 #####################################################
@@ -75,6 +78,8 @@ def predict_pipeline_pytorch(
     keypt_color,
     marker_size,
     flag_color_by_confidence,
+    colormap,
+    bbox_color,
 ):
     # detection + pose with the SuperAnimal PyTorch models (keypoints in image coords)
     img_output, animals, bodyparts = predict_superanimal(
@@ -94,16 +99,22 @@ def predict_pipeline_pytorch(
             keypt_color=keypt_color,
             marker_size=marker_size,
             color_by_confidence=flag_color_by_confidence,
+            colormap=colormap,
         )
         if not flag_dlc_only:
-            draw_bbox_w_text(img_output, animal["bbox"], font_size=font_size)
+            draw_bbox_w_text(img_output, animal["bbox"], font_size=font_size, bbox_color=bbox_color)
 
     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 finalize_outputs(
-        img_output, download_file, [animal["kpts"] for animal in animals], map_label_id_to_str, flag_color_by_confidence
+        img_output,
+        download_file,
+        [animal["kpts"] for animal in animals],
+        map_label_id_to_str,
+        flag_color_by_confidence,
+        colormap,
     )
 
 
@@ -122,6 +133,8 @@ def predict_pipeline(
     keypt_color,
     marker_size,
     flag_color_by_confidence,
+    colormap,
+    bbox_color,
 ):
 
     if backend == "PyTorch":
@@ -137,6 +150,8 @@ def predict_pipeline(
             keypt_color,
             marker_size,
             flag_color_by_confidence,
+            colormap,
+            bbox_color,
         )
 
     # TensorFlow (legacy): MegaDetector crops + DLCLive
@@ -202,12 +217,13 @@ def predict_pipeline(
             keypt_color=keypt_color,
             marker_size=marker_size,
             color_by_confidence=flag_color_by_confidence,
+            colormap=colormap,
         )
 
         donw_file = save_results_only_dlc(list_kpts_per_crop[0], map_label_id_to_str, dlc_model_name)
 
         return finalize_outputs(
-            img_input, donw_file, [list_kpts_per_crop[0]], map_label_id_to_str, flag_color_by_confidence
+            img_input, donw_file, [list_kpts_per_crop[0]], map_label_id_to_str, flag_color_by_confidence, colormap
         )
 
     else:
@@ -233,13 +249,14 @@ def predict_pipeline(
                 keypt_color=keypt_color,
                 marker_size=marker_size,
                 color_by_confidence=flag_color_by_confidence,
+                colormap=colormap,
             )
 
             # Paste crop in original image
             img_background.paste(img_crop, box=tuple([int(t) for t in bb_per_animal[:2]]))
 
             # Plot bbox
-            draw_bbox_w_text(img_background, bb_per_animal, font_size=font_size)  # TODO: add selectable color for bbox?
+            draw_bbox_w_text(img_background, bb_per_animal, font_size=font_size, bbox_color=bbox_color)
 
         # Save detection results as json
         download_file = save_results_as_json(
@@ -247,7 +264,7 @@ def predict_pipeline(
         )
 
         return finalize_outputs(
-            img_background, download_file, list_kpts_per_crop, map_label_id_to_str, flag_color_by_confidence
+            img_background, download_file, list_kpts_per_crop, map_label_id_to_str, flag_color_by_confidence, colormap
         )
 
 
@@ -280,4 +297,4 @@ with gr.Blocks(title=gr_title) as demo:
 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())
+demo.launch(theme=dlc_theme())

ui_utils.py

@@ -1,4 +1,54 @@
 import gradio as gr
+from matplotlib import colormaps
+from matplotlib.colors import to_hex
+from PIL import Image
+
+from pytorch_utils import MAX_IMAGE_SIZE
+from viz_utils import COLORMAPS
+
+# shades built around the DeepLabCut docs palette (DeepLabCut/docs/_static/custom.css)
+DLC_PURPLE = gr.themes.Color(
+    name="dlc_purple",
+    c50="#f5edff",
+    c100="#ead7ff",
+    c200="#d9b8ff",
+    c300="#c084fc",
+    c400="#ac72f0",
+    c500="#9b5de5",
+    c600="#8550c4",
+    c700="#73439a",
+    c800="#5c357c",
+    c900="#4b236f",
+    c950="#2e1546",
+)
+DLC_TEAL = gr.themes.Color(
+    name="dlc_teal",
+    c50="#e8f8f7",
+    c100="#c9efec",
+    c200="#9fe2dd",
+    c300="#7ad3ce",
+    c400="#57c4be",
+    c500="#2fb0a8",
+    c600="#21a197",
+    c700="#1a8078",
+    c800="#16655f",
+    c900="#124f4a",
+    c950="#0a2e2b",
+)
+
+
+def dlc_theme():
+    # white text on purple 500 is 4.1:1, so filled buttons use 700 (7.0:1) and 600 (5.3:1)
+    return gr.themes.Default(primary_hue=DLC_PURPLE, secondary_hue=DLC_TEAL, neutral_hue="slate").set(
+        button_primary_background_fill="*primary_700",
+        button_primary_background_fill_hover="*primary_600",
+        button_primary_background_fill_dark="*primary_600",
+        button_primary_background_fill_hover_dark="*primary_500",
+        button_primary_text_color="white",
+        button_primary_text_color_dark="white",
+        button_primary_border_color="*primary_700",
+        button_primary_border_color_dark="*primary_600",
+    )
 
 
 def gradio_inputs_for_MD_DLC(backends_list, md_models_list, dlc_models_list):
@@ -62,11 +112,23 @@ def gradio_inputs_for_MD_DLC(backends_list, md_models_list, dlc_models_list):
             label="Color keypoints by confidence? (otherwise by bodypart)",
         )
 
+        gr_colormap = gr.Dropdown(
+            choices=COLORMAPS,
+            value="viridis",
+            type="value",
+            label="Keypoint colormap",
+        )
+
         gr_keypt_color = gr.ColorPicker(
             value="#862db7",
             label="Choose color for keypoint label",
         )
 
+        gr_bbox_color = gr.ColorPicker(
+            value="#ff0000",
+            label="Choose color for bounding boxes",
+        )
+
         gr_labels_font_style = gr.Dropdown(
             choices=["amiko", "animals", "nature", "painter", "zen"],
             value="amiko",
@@ -77,7 +139,7 @@ def gradio_inputs_for_MD_DLC(backends_list, md_models_list, dlc_models_list):
         gr_slider_font_size = gr.Slider(
             minimum=5,
             maximum=30,
-            value=8,
+            value=18,
             step=1,
             label="Set font size",
         )
@@ -85,7 +147,7 @@ def gradio_inputs_for_MD_DLC(backends_list, md_models_list, dlc_models_list):
         gr_slider_marker_size = gr.Slider(
             minimum=1,
             maximum=20,
-            value=9,
+            value=6,
             step=1,
             label="Set marker size",
         )
@@ -104,11 +166,29 @@ def gradio_inputs_for_MD_DLC(backends_list, md_models_list, dlc_models_list):
         gr_keypt_color,
         gr_slider_marker_size,
         gr_color_by_confidence_checkbox,
+        gr_colormap,
+        gr_bbox_color,
     ]
 
 
+def confidence_legend_html(colormap="viridis"):
+    # the colormap as a CSS gradient, matching the keypoint fill in draw_keypoints_on_image
+    gradient = ", ".join(f"{to_hex(colormaps[colormap](i / 10))} {i * 10}%" for i in range(11))
+    # no leading newline: gradio prefixes "'" to cached example values starting with one (CSV injection guard)
+    return f"""<div style="display:flex; align-items:flex-start; gap:12px; flex-wrap:wrap; font-size:var(--text-sm);">
+  <span style="white-space:nowrap; line-height:12px;">Keypoint confidence</span>
+  <div style="flex:1; min-width:160px; max-width:360px;">
+    <div style="height:12px; border-radius:6px; background:linear-gradient(to right, {gradient});"></div>
+    <div style="display:flex; justify-content:space-between; margin-top:2px; font-variant-numeric:tabular-nums;">
+      <span>0</span><span>0.5</span><span>1</span>
+    </div>
+  </div>
+</div>"""
+
+
 def gradio_outputs_for_MD_DLC():
     gr_image_output = gr.Image(type="pil", label="Output Image")
+    gr_confidence_legend = gr.HTML("")
     with gr.Row():
         gr_file_download = gr.File(label="Download JSON file")
         gr_image_download = gr.File(label="Download annotated image")
@@ -117,7 +197,13 @@ def gradio_outputs_for_MD_DLC():
         label="Keypoint confidence (lowest first)",
         interactive=False,
     )
-    return [gr_image_output, gr_file_download, gr_image_download, gr_confidence_table]
+    return [gr_image_output, gr_confidence_legend, gr_file_download, gr_image_download, gr_confidence_table]
+
+
+def example_sizes(path):
+    # font and marker sizes for the resolution the PyTorch backend draws on
+    side = min(max(Image.open(path).size), MAX_IMAGE_SIZE)
+    return round(side / 64), round(side / 200)
 
 
 def gradio_description_and_examples():
@@ -130,8 +216,16 @@ def gradio_description_and_examples():
     )
 
     examples = [
-        [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")
+        [image, "PyTorch", "md_v5a", "superanimal_quadruped", False, True, 0.5, 0.4, "amiko"]
+        + [font_size, "#ff0000", marker_size, True, "viridis", "#ff0000"]
+        for image in (
+            "examples/dog.jpeg",
+            "examples/cat.jpg",
+            "examples/lynx.jpg",
+            "examples/goat.jpg",
+            "examples/giraffe.jpg",
+        )
+        for font_size, marker_size in [example_sizes(image)]
     ]
 
     return [title, description, examples]

viz_utils.py

@@ -2,8 +2,8 @@ import json
 from datetime import date
 
 import numpy as np
-from matplotlib import cm
-from PIL import Image, ImageColor, ImageDraw, ImageFont
+from matplotlib import colormaps
+from PIL import ImageColor, ImageDraw, ImageFont
 
 today = date.today()
 FONTS = {
@@ -13,6 +13,8 @@ FONTS = {
     "animals": "fonts/UncialAnimals.ttf",
     "zen": "fonts/ZEN.TTF",
 }
+# perceptually uniform maps first; turbo separates neighbouring bodyparts best
+COLORMAPS = ["viridis", "plasma", "magma", "cividis", "turbo"]
 
 
 #########################################
@@ -28,6 +30,7 @@ def draw_keypoints_on_image(
     keypt_color="#ff0000",
     marker_size=2,
     color_by_confidence=True,
+    colormap="viridis",
 ):
     """Draws keypoints on an image.
     Modified from:
@@ -57,7 +60,7 @@ def draw_keypoints_on_image(
         keypoints_x = tuple([im_width * x for x in keypoints_x])
         keypoints_y = tuple([im_height * y for y in keypoints_y])
 
-    # cmap = matplotlib.cm.get_cmap('hsv')
+    cmap = colormaps[colormap]
     # draw ellipses around keypoints
     for i, (keypoint_x, keypoint_y) in enumerate(zip(keypoints_x, keypoints_y, strict=True)):
         # handling potential nans in the keypoints
@@ -66,11 +69,11 @@ def draw_keypoints_on_image(
 
         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)
+            # fill color encodes the keypoint confidence (see confidence_legend_html in ui_utils)
+            round_fill = cmap(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 = list(cmap(i / max(len(keypoints) - 1, 1), bytes=True))
             round_fill[3] = round(confidence * 255)
             round_fill = tuple(round_fill)
         draw.ellipse(
@@ -88,39 +91,29 @@ def draw_keypoints_on_image(
             font = ImageFont.truetype(FONTS[font_style], font_size)
             draw.text(
                 (keypoint_x + marker_size, keypoint_y + marker_size),  # (0.5*im_width, 0.5*im_height), #-------
-                map_label_id_to_str[i],
+                display_bodypart(map_label_id_to_str[i]),
                 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).
+#  Bodypart names for display
+# display names where the SuperAnimal definitions misspell (quadruped "thai") or read oddly
+# (top-view mouse "backend"); the JSON output keeps the model's names
+BODYPART_DISPLAY_NAMES = {
+    "front_left_thai": "front left thigh",
+    "front_right_thai": "front right thigh",
+    "back_left_thai": "back left thigh",
+    "back_right_thai": "back right thigh",
+    "mid_backend": "mid back end",
+    "mid_backend2": "mid back end 2",
+    "mid_backend3": "mid back end 3",
+}
 
-    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
+
+def display_bodypart(name):
+    return BODYPART_DISPLAY_NAMES.get(name, name.replace("_", " "))
 
 
 #########################################
@@ -131,7 +124,7 @@ def keypoint_confidence_rows(kpts_per_animal, map_label_id_to_str):
     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)])
+                rows.append([i_animal, display_bodypart(map_label_id_to_str[i_kpt]), round(float(kpt[2]), 3)])
     return sorted(rows, key=lambda row: row[2])
 
 
@@ -144,25 +137,27 @@ def save_annotated_image(image, path_to_output_file="download_annotated.png"):
 
 #########################################
 #  Draw bboxes on image
-def draw_bbox_w_text(img, results, font_style="amiko", font_size=8):  # TODO: select color too?
-    # pdb.set_trace()
-    bbxyxy = results
-    w, h = bbxyxy[2], bbxyxy[3]
-    shape = [(bbxyxy[0], bbxyxy[1]), (w, h)]
-    imgR = ImageDraw.Draw(img)
-    imgR.rectangle(shape, outline="red", width=5)  ##bb for animal
-
-    confidence = bbxyxy[4]
-    string_bb = "animal " + str(round(confidence, 2))
-    font = ImageFont.truetype(FONTS[font_style], font_size)
-
-    text_size = font.getbbox(string_bb)  # (h,w)
-    position = (bbxyxy[0], bbxyxy[1] - text_size[1] - 2)
-    left, top, right, bottom = imgR.textbbox(position, string_bb, font=font)
-    imgR.rectangle((left, top - 5, right + 5, bottom + 5), fill="red")
-    imgR.text((bbxyxy[0] + 3, bbxyxy[1] - text_size[1] - 2), string_bb, font=font, fill="black")
+def draw_bbox_w_text(img, results, font_style="amiko", font_size=8, bbox_color="#ff0000"):
+    x1, y1, x2, y2, confidence = results[:5]
+    draw = ImageDraw.Draw(img)
+    draw.rectangle([(x1, y1), (x2, y2)], outline=bbox_color, width=max(2, round(font_size / 5)))
 
-    return imgR
+    label = f"animal {confidence:.2f}"
+    font = ImageFont.truetype(FONTS[font_style], font_size)
+    left, top, right, bottom = draw.textbbox((0, 0), label, font=font)
+    pad = max(2, font_size // 5)
+    label_w, label_h = right - left + 2 * pad, bottom - top + 2 * pad
+    # label above the box, or inside it when the box touches the top of the image
+    label_y = y1 - label_h if y1 >= label_h else y1
+    draw.rectangle([(x1, label_y), (x1 + label_w, label_y + label_h)], fill=bbox_color)
+    draw.text((x1 + pad - left, label_y + pad - top), label, font=font, fill=label_text_color(bbox_color))
+
+
+def label_text_color(background):
+    # black or white, whichever contrasts more with the background (WCAG relative luminance)
+    channels = [c / 255 for c in ImageColor.getrgb(background)[:3]]
+    r, g, b = [c / 12.92 if c <= 0.03928 else ((c + 0.055) / 1.055) ** 2.4 for c in channels]
+    return "black" if 0.2126 * r + 0.7152 * g + 0.0722 * b > 0.179 else "white"
 
 
 ###########################################