Add PyTorch SuperAnimal backend and refresh the Space

.gitignore

@@ -1,3 +1,16 @@
 # Byte-compiled / optimized / DLL files
 __pycache__/
-model/__pycache__/
+model/__pycache__/
+.vscode/settings.json
+
+# Local environment and tool caches
+.venv/
+.ruff_cache/
+
+# Written by the app at runtime
+.gradio/
+download_predictions.json
+dowload_predictions_dlc.json
+download_annotated.png
+# TF (legacy) models downloaded on first use; DLC_models/readme.md stays tracked
+DLC_models/*/

.pre-commit-config.yaml

@@ -0,0 +1,28 @@
+# Adapted from DeepLabCut's .pre-commit-config.yaml.
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v6.0.0
+    hooks:
+      - id: check-added-large-files
+      - id: check-yaml
+      - id: check-toml
+      - id: check-merge-conflict
+      - id: end-of-file-fixer
+      - id: trailing-whitespace
+
+  - repo: https://github.com/tox-dev/pyproject-fmt
+    rev: v2.19.0
+    hooks:
+      - id: pyproject-fmt
+
+  - repo: https://github.com/abravalheri/validate-pyproject
+    rev: v0.25
+    hooks:
+      - id: validate-pyproject
+
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.15.6
+    hooks:
+      - id: ruff-check
+        args: [--fix, --unsafe-fixes]
+      - id: ruff-format

README.md

@@ -5,9 +5,9 @@ colorFrom: blue
 colorTo: purple
 sdk: gradio
 python_version: 3.12
-sdk_version: 6.11.0
+sdk_version: 6.29.0
 app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,193 +1,314 @@
-# Adapted from https://huggingface.co/spaces/hlydecker/MegaDetector_v5 
+# Adapted from https://huggingface.co/spaces/hlydecker/MegaDetector_v5
 # Adapted from https://huggingface.co/spaces/sofmi/MegaDetector_DLClive/blob/main/app.py
-# Adapted from https://huggingface.co/spaces/Neslihan/megadetector_dlcmodels/blob/main/app.py 
+# Adapted from https://huggingface.co/spaces/Neslihan/megadetector_dlcmodels/blob/main/app.py
 # Adapted from  https://huggingface.co/spaces/DeepLabCut/MegaDetector_DeepLabCut
 
 import os
-import yaml
-import numpy as np
-from matplotlib import cm
-import gradio as gr
-import deeplabcut
-import dlclibrary
-import dlclive
-# import transformers
-
-from PIL import Image, ImageColor, ImageFont, ImageDraw 
-import requests
+import threading
 
-from viz_utils import save_results_as_json, draw_keypoints_on_image, draw_bbox_w_text, save_results_only_dlc
-from detection_utils import predict_md, crop_animal_detections
-from dlc_utils import predict_dlc
-from ui_utils import gradio_inputs_for_MD_DLC, gradio_outputs_for_MD_DLC, gradio_description_and_examples
-
-from deeplabcut.utils import auxiliaryfunctions
+import gradio as gr
+import numpy as np
+import yaml
 from dlclibrary.dlcmodelzoo.modelzoo_download import (
     download_huggingface_model,
-    MODELOPTIONS,
 )
-from dlclive import DLCLive, Processor
+from dlclive import Processor
 
+# import transformers
+from PIL import Image
 
-# TESTING (passes) download the SuperAnimal models:
-#model = 'superanimal_topviewmouse'
-#train_dir = 'DLC_models/sa-tvm'
-#download_huggingface_model(model, train_dir)
+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 (
+    confidence_legend_html,
+    dlc_theme,
+    gradio_description_and_examples,
+    gradio_inputs_for_MD_DLC,
+    gradio_outputs_for_MD_DLC,
+)
+from viz_utils import (
+    draw_bbox_w_text,
+    draw_keypoints_on_image,
+    keypoint_confidence_rows,
+    save_annotated_image,
+    save_results_as_json,
+    save_results_only_dlc,
+    save_results_pytorch,
+)
 
-# grab demo data cooco cat:
-url = "http://images.cocodataset.org/val2017/000000039769.jpg"
-image = Image.open(requests.get(url, stream=True).raw)
+# TESTING (passes) download the SuperAnimal models:
+# model = 'superanimal_topviewmouse'
+# train_dir = 'DLC_models/sa-tvm'
+# download_huggingface_model(model, train_dir)
 
 # megadetector and dlc model look up
-MD_models_dict = {'md_v5a': "MD_models/md_v5a.0.0.pt", # 
-                  'md_v5b': "MD_models/md_v5b.0.0.pt"}
+MD_models_dict = {
+    "md_v5a": "MD_models/md_v5a.0.0.pt",  #
+    "md_v5b": "MD_models/md_v5b.0.0.pt",
+}
 
-# DLC models target  dirs
-DLC_models_dict = {'superanimal_topviewmouse_dlcrnet': "DLC_models/sa-tvm",
-                   'superanimal_quadruped_dlcrnet': "DLC_models/sa-q"}
-                    
+BACKENDS = ["PyTorch", "TensorFlow (legacy)"]
+
+# TF (legacy) DLC models: model zoo name and target dir, per SuperAnimal
+DLC_models_dict = {
+    "superanimal_topviewmouse": ("superanimal_topviewmouse_dlcrnet", "DLC_models/sa-tvm"),
+    "superanimal_quadruped": ("superanimal_quadruped_dlcrnet", "DLC_models/sa-q"),
+}
 
 
 #####################################################
-def predict_pipeline(img_input,
-                     mega_model_input,
-                     dlc_model_input_str,
-                     flag_dlc_only,
-                     flag_show_str_labels,
-                     bbox_likelihood_th,
-                     kpts_likelihood_th,
-                     font_style,
-                     font_size,
-                     keypt_color,
-                     marker_size,
-                     ):
+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)
+    legend = confidence_legend_html(colormap) if color_by_confidence else ""
+    return img_output, legend, download_file, annotated_file, confidence_rows
+
+
+#####################################################
+def predict_pipeline_pytorch(
+    img_input,
+    superanimal,
+    flag_dlc_only,
+    flag_show_str_labels,
+    bbox_likelihood_th,
+    kpts_likelihood_th,
+    font_style,
+    font_size,
+    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(
+        img_input, superanimal, bbox_likelihood_th, kpts_likelihood_th, full_image=flag_dlc_only
+    )
+    map_label_id_to_str = dict(enumerate(bodyparts))
+
+    for animal in animals:
+        draw_keypoints_on_image(
+            img_output,
+            animal["kpts"],
+            map_label_id_to_str,
+            flag_show_str_labels,
+            use_normalized_coordinates=False,
+            font_style=font_style,
+            font_size=font_size,
+            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, 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,
+        image_size=img_input.size,
+        annotated_size=img_output.size,
+    )
+    return finalize_outputs(
+        img_output,
+        download_file,
+        [animal["kpts"] for animal in animals],
+        map_label_id_to_str,
+        flag_color_by_confidence,
+        colormap,
+    )
+
+
+#####################################################
+def predict_pipeline(
+    img_input,
+    backend,
+    mega_model_input,
+    dlc_model_input_str,
+    flag_dlc_only,
+    flag_show_str_labels,
+    bbox_likelihood_th,
+    kpts_likelihood_th,
+    font_style,
+    font_size,
+    keypt_color,
+    marker_size,
+    flag_color_by_confidence,
+    colormap,
+    bbox_color,
+):
+
+    if backend == "PyTorch":
+        return predict_pipeline_pytorch(
+            img_input,
+            dlc_model_input_str,
+            flag_dlc_only,
+            flag_show_str_labels,
+            bbox_likelihood_th,
+            kpts_likelihood_th,
+            font_style,
+            font_size,
+            keypt_color,
+            marker_size,
+            flag_color_by_confidence,
+            colormap,
+            bbox_color,
+        )
+
+    # TensorFlow (legacy): MegaDetector crops + DLCLive
+    dlc_model_name, dlc_model_dir = DLC_models_dict[dlc_model_input_str]
 
     if not flag_dlc_only:
-        ############################################################                                               
+        ############################################################
         # ### Run Megadetector
-        md_results = predict_md(img_input, 
-                                MD_models_dict[mega_model_input], #mega_model_input,
-                                size=640) #Image.fromarray(results.imgs[0])
+        md_results = predict_md(
+            img_input,
+            MD_models_dict[mega_model_input],  # mega_model_input,
+            size=640,
+        )  # Image.fromarray(results.imgs[0])
 
         ################################################################
-        # Obtain animal crops for bboxes with confidence above th
-        list_crops = crop_animal_detections(img_input,
-                                            md_results,
-                                            bbox_likelihood_th)
+        # Obtain animal crops (and their bboxes) with confidence above th
+        list_crops, list_bboxes = crop_animal_detections(img_input, md_results, bbox_likelihood_th)
 
         ############################################################
 
-    ## Get DLC model and label map  
-    
+    ## Get DLC model and label map
+
     # If model is found: do not download (previous execution is likely within same day)
     # TODO: can we ask the user whether to reload dlc model if a directory is found?
-    if os.path.isdir(DLC_models_dict[dlc_model_input_str]) and \
-        len(os.listdir(DLC_models_dict[dlc_model_input_str])) > 0:
-        path_to_DLCmodel = DLC_models_dict[dlc_model_input_str]
-    else:
-        path_to_DLCmodel = DLC_models_dict[dlc_model_input_str]
-        download_huggingface_model(dlc_model_input_str, path_to_DLCmodel)
+    path_to_DLCmodel = dlc_model_dir
+    if not (os.path.isdir(dlc_model_dir) and len(os.listdir(dlc_model_dir)) > 0):
+        download_huggingface_model(dlc_model_name, path_to_DLCmodel)
 
     # extract map label ids to strings
-    pose_cfg_path = os.path.join(DLC_models_dict[dlc_model_input_str],
-                                 'pose_cfg.yaml')
-    with open(pose_cfg_path, "r") as stream:
-        pose_cfg_dict = yaml.safe_load(stream) 
-    map_label_id_to_str = dict([(k,v) for k,v in zip([el[0] for el in pose_cfg_dict['all_joints']],  # pose_cfg_dict['all_joints'] is a list of one-element lists,
-                                                     pose_cfg_dict['all_joints_names'])])
-
-
-##############################################################
+    pose_cfg_path = os.path.join(dlc_model_dir, "pose_cfg.yaml")
+    with open(pose_cfg_path) as stream:
+        pose_cfg_dict = yaml.safe_load(stream)
+    map_label_id_to_str = dict(
+        [
+            (k, v)
+            for k, v in zip(
+                [
+                    el[0] for el in pose_cfg_dict["all_joints"]
+                ],  # pose_cfg_dict['all_joints'] is a list of one-element lists,
+                pose_cfg_dict["all_joints_names"],
+                strict=True,
+            )
+        ]
+    )
+
+    ##############################################################
     # Run DLC and visualize results
-    dlc_proc = Processor() #TODO: update deeplabcut.video_inference_superanimal() once merged
+    dlc_proc = Processor()  # TODO: update deeplabcut.video_inference_superanimal() once merged
 
     # if required: ignore MD crops and run DLC on full image [mostly for testing]
     if flag_dlc_only:
         # compute kpts on input img
-        list_kpts_per_crop = predict_dlc([np.asarray(img_input)],
-                                         kpts_likelihood_th,
-                                         path_to_DLCmodel,
-                                         dlc_proc)
+        list_kpts_per_crop = predict_dlc([np.asarray(img_input)], kpts_likelihood_th, path_to_DLCmodel, dlc_proc)
         # draw kpts on input img #fix!
-        draw_keypoints_on_image(img_input,
-                                list_kpts_per_crop[0], # a numpy array with shape [num_keypoints, 2].
-                                map_label_id_to_str,
-                                flag_show_str_labels,
-                                use_normalized_coordinates=False,
-                                font_style=font_style,
-                                font_size=font_size,
-                                keypt_color=keypt_color,
-                                marker_size=marker_size)
-
-        donw_file = save_results_only_dlc(list_kpts_per_crop[0], map_label_id_to_str,dlc_model_input_str)
-
-        return img_input, donw_file
+        draw_keypoints_on_image(
+            img_input,
+            list_kpts_per_crop[0],  # a numpy array with shape [num_keypoints, 2].
+            map_label_id_to_str,
+            flag_show_str_labels,
+            use_normalized_coordinates=False,
+            font_style=font_style,
+            font_size=font_size,
+            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, image_size=img_input.size
+        )
+
+        return finalize_outputs(
+            img_input, donw_file, [list_kpts_per_crop[0]], map_label_id_to_str, flag_color_by_confidence, colormap
+        )
 
     else:
         # Compute kpts for each crop
-        list_kpts_per_crop = predict_dlc(list_crops,
-                                         kpts_likelihood_th,
-                                         path_to_DLCmodel,
-                                         dlc_proc)
-        
+        list_kpts_per_crop = predict_dlc(list_crops, kpts_likelihood_th, path_to_DLCmodel, dlc_proc)
+
         # resize input image to match megadetector output
-        img_background = img_input.resize((md_results.ims[0].shape[1],
-                                           md_results.ims[0].shape[0]))
-        
-        # draw keypoints on each crop and paste to background img
-        for ic, (np_crop, kpts_crop) in enumerate(zip(list_crops,
-                                                      list_kpts_per_crop)):
+        img_background = img_input.resize((md_results.ims[0].shape[1], md_results.ims[0].shape[0]))
 
+        # draw keypoints on each crop and paste to background img
+        for np_crop, kpts_crop, bb_per_animal in zip(list_crops, list_kpts_per_crop, list_bboxes, strict=True):
             img_crop = Image.fromarray(np_crop)
 
             # Draw keypts on crop
-            draw_keypoints_on_image(img_crop,
-                                    kpts_crop, # a numpy array with shape [num_keypoints, 2].
-                                    map_label_id_to_str,
-                                    flag_show_str_labels,
-                                    use_normalized_coordinates=False,  # if True, then I should use md_results.xyxyn for list_kpts_crop
-                                    font_style=font_style,
-                                    font_size=font_size,
-                                    keypt_color=keypt_color,
-                                    marker_size=marker_size)
+            draw_keypoints_on_image(
+                img_crop,
+                kpts_crop,  # a numpy array with shape [num_keypoints, 2].
+                map_label_id_to_str,
+                flag_show_str_labels,
+                use_normalized_coordinates=False,  # if True, then I should use md_results.xyxyn for list_kpts_crop
+                font_style=font_style,
+                font_size=font_size,
+                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 md_results.xyxy[0][ic,:2]]))
+            img_background.paste(img_crop, box=tuple([int(t) for t in bb_per_animal[:2]]))
 
             # Plot bbox
-            bb_per_animal = md_results.xyxy[0].tolist()[ic]
-            pred = md_results.xyxy[0].tolist()[ic][4]
-            if bbox_likelihood_th < pred:
-                draw_bbox_w_text(img_background, 
-                                    bb_per_animal,
-                                    font_size=font_size)  # TODO: add selectable color for bbox?
-
-            
-        # Save detection results as json
-        download_file  = save_results_as_json(md_results,list_kpts_per_crop,map_label_id_to_str, bbox_likelihood_th,dlc_model_input_str,mega_model_input)         
+            draw_bbox_w_text(img_background, bb_per_animal, font_size=font_size, bbox_color=bbox_color)
 
-        return img_background, download_file
+        # 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,
+            image_size=img_input.size,
+        )
 
+        return finalize_outputs(
+            img_background, download_file, list_kpts_per_crop, map_label_id_to_str, flag_color_by_confidence, colormap
+        )
 
 
 #########################################################
 # Define user interface and launch
-inputs = gradio_inputs_for_MD_DLC(list(MD_models_dict.keys()),
-                                  list(DLC_models_dict.keys()))
-outputs = gradio_outputs_for_MD_DLC()                                    
-[gr_title, 
- gr_description, 
- examples] = gradio_description_and_examples()
-
-# launch
-demo = gr.Interface(predict_pipeline, 
-                    inputs=inputs,
-                    outputs=outputs, 
-                    title=gr_title, 
-                    description=gr_description,
-                    examples=examples,
-                    )
-
-demo.queue()
-demo.launch(theme=gr.themes.Default())
+[gr_title, gr_description, examples] = gradio_description_and_examples()
+
+with gr.Blocks(title=gr_title) as demo:
+    gr.Markdown(f"# {gr_title}\n{gr_description}")
+    with gr.Row():
+        with gr.Column():
+            inputs = gradio_inputs_for_MD_DLC(BACKENDS, list(MD_models_dict.keys()), list(DLC_models_dict.keys()))
+            run_button = gr.Button("Run", variant="primary")
+        with gr.Column():
+            outputs = gradio_outputs_for_MD_DLC()
+
+    # the MegaDetector choice only applies to the TensorFlow (legacy) backend
+    gr_backend_input, gr_mega_model_input = inputs[1], inputs[2]
+    gr_backend_input.change(
+        lambda backend: gr.update(visible=backend != "PyTorch"), inputs=gr_backend_input, outputs=gr_mega_model_input
+    )
+
+    run_button.click(predict_pipeline, inputs=inputs, outputs=outputs, api_name="predict")
+
+    # cached on first click, so a failing download cannot block startup
+    gr.Examples(examples, inputs=inputs, outputs=outputs, fn=predict_pipeline, 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=dlc_theme())

detection_utils.py

@@ -1,92 +1,74 @@
-
-from tkinter import W
-import gradio as gr
-from matplotlib import cm
-import torch
-import torchvision
-import matplotlib
-import PIL
-from PIL import Image, ImageColor, ImageFont, ImageDraw 
-import numpy as np
 import math
 
+import numpy as np
+import PIL
+import torch
 
-import yaml
-import pdb
 
 ############################################
 # Predict detections with MegaDetector v5a model
-def predict_md(im, 
-               megadetector_model, #Megadet_Models[mega_model_input]
-               size=640):
-               
+def predict_md(
+    im,
+    megadetector_model,  # Megadet_Models[mega_model_input]
+    size=640,
+):
+
     # resize image
-    g = (size / max(im.size))  # multipl factor to make max size of the image equal to input size
-    im = im.resize((int(x * g) for x in im.size), 
-                    PIL.Image.Resampling.LANCZOS)  # resize
-    # device
-    if torch.cuda.is_available():
-        md_device = torch.device('cuda')
-    else:
-        md_device = torch.device('cpu')
-
-    # megadetector    
-    MD_model = torch.hub.load('ultralytics/yolov5', # repo_or_dir
-                              'custom', #model
-                              megadetector_model, # args for callable model
-                              skip_validation=True, # avoid GitHub API rate limit (403)
-                              device=md_device,
-                              trust_repo=True
-                              ) 
-                              
-    # send model to gpu if possible
-    if (md_device == torch.device('cuda')):
-        print('Sending model to GPU')
-        MD_model.to(md_device)    
+    g = size / max(im.size)  # multipl factor to make max size of the image equal to input size
+    im = im.resize((int(x * g) for x in im.size), PIL.Image.Resampling.LANCZOS)  # resize
+    # device: yolov5's select_device expects a CUDA index ('0') or 'cpu', not 'cuda'
+    md_device = "0" if torch.cuda.is_available() else "cpu"
+
+    # megadetector
+    MD_model = torch.hub.load(
+        "ultralytics/yolov5",  # repo_or_dir
+        "custom",  # model
+        megadetector_model,  # args for callable model
+        skip_validation=True,  # avoid GitHub API rate limit (403)
+        device=md_device,
+        trust_repo=True,
+    )
 
     ## detect objects
-    results = MD_model(im)  # inference # vars(results).keys()= dict_keys(['imgs', 'pred', 'names', 'files', 'times', 'xyxy', 'xywh', 'xyxyn', 'xywhn', 'n', 't', 's'])
-    
-    return results  
+    # vars(results).keys(): imgs, pred, names, files, times, xyxy, xywh, xyxyn, xywhn, n, t, s
+    results = MD_model(im)
+
+    return results
 
 
 ##########################################
-def crop_animal_detections(img_in,
-                           yolo_results, 
-                           likelihood_th):
+def crop_animal_detections(img_in, yolo_results, likelihood_th):
 
     ## Extract animal crops
-    list_labels_as_str = [i for i in yolo_results.names.values()]  # ['animal', 'person', 'vehicle'] 
+    list_labels_as_str = [i for i in yolo_results.names.values()]  # ['animal', 'person', 'vehicle']
     list_np_animal_crops = []
+    list_animal_bboxes = []  # detection rows [x1,y1,x2,y2,conf,label] matching each crop
 
     # image to crop (scale as input for megadetector)
-    img_in = img_in.resize((yolo_results.ims[0].shape[1],
-                            yolo_results.ims[0].shape[0]))
-    # for every detection in the img                        
+    img_in = img_in.resize((yolo_results.ims[0].shape[1], yolo_results.ims[0].shape[0]))
+    # for every detection in the img
     for det_array in yolo_results.xyxy:
-
         # for every detection
         for j in range(det_array.shape[0]):
-
             # compute coords around bbox rounded to the nearest integer (for pasting later)
-            xmin_rd = int(math.floor(det_array[j,0])) # int() should suffice?
-            ymin_rd = int(math.floor(det_array[j,1]))
+            xmin_rd = int(math.floor(det_array[j, 0]))  # int() should suffice?
+            ymin_rd = int(math.floor(det_array[j, 1]))
 
-            xmax_rd = int(math.ceil(det_array[j,2]))
-            ymax_rd = int(math.ceil(det_array[j,3]))
+            xmax_rd = int(math.ceil(det_array[j, 2]))
+            ymax_rd = int(math.ceil(det_array[j, 3]))
 
-            pred_llk = det_array[j,4] 
-            pred_label = det_array[j,5]
+            pred_llk = det_array[j, 4]
+            pred_label = det_array[j, 5]
             # keep animal crops above threshold
-            if (pred_label == list_labels_as_str.index('animal')) and \
-                (pred_llk >= likelihood_th):
+            if (pred_label == list_labels_as_str.index("animal")) and (pred_llk >= likelihood_th):
                 area = (xmin_rd, ymin_rd, xmax_rd, ymax_rd)
 
-                #pdb.set_trace()
-                crop = img_in.crop(area) #Image.fromarray(img_in).crop(area)
+                # pdb.set_trace()
+                crop = img_in.crop(area)  # Image.fromarray(img_in).crop(area)
                 crop_np = np.asarray(crop)
 
                 # add to list
                 list_np_animal_crops.append(crop_np)
+                list_animal_bboxes.append(det_array[j, :].tolist())
 
-    return list_np_animal_crops
+    return list_np_animal_crops, list_animal_bboxes

dlc_utils.py

@@ -1,32 +1,28 @@
-import deeplabcut
-from tkinter import W
-import gradio as gr
 import numpy as np
-from dlclive import DLCLive, Processor
+from dlclive import DLCLive
 
 
 ##########################################
-def predict_dlc(list_np_crops,
-                kpts_likelihood_th,
-                dlc_model_folder,
-                dlc_proc):
-    
+def predict_dlc(list_np_crops, kpts_likelihood_th, dlc_model_folder, dlc_proc):
+
+    # no animal detected: nothing to run
+    if len(list_np_crops) == 0:
+        return []
+
+    # DLCLive always converts 3-channel frames BGR->RGB, but our crops are
+    # already RGB (PIL), so pass them as BGR for the model to see RGB
+    list_np_crops = [np.ascontiguousarray(crop[..., ::-1]) for crop in list_np_crops]
+
     # run dlc thru list of crops
     dlc_live = DLCLive(dlc_model_folder, processor=dlc_proc)
     dlc_live.init_inference(list_np_crops[0])
 
     list_kpts_per_crop = []
-    all_kypts  = []
-    np_aux = np.empty((1,3)) # can I avoid hardcoding here?
     for crop in list_np_crops:
-        # scale crop here?
-        keypts_xyp = dlc_live.get_pose(crop) # third column is llk!
+        keypts_xyp = dlc_live.get_pose(crop)  # third column is llk!
         # set kpts below threhsold to nan
-        
-        #pdb.set_trace()
-        keypts_xyp[keypts_xyp[:,-1] < kpts_likelihood_th,:] = np_aux.fill(np.nan)
-        # add kpts of this crop to list 
+        keypts_xyp[keypts_xyp[:, -1] < kpts_likelihood_th, :] = np.nan
+        # add kpts of this crop to list
         list_kpts_per_crop.append(keypts_xyp)
-        all_kypts.append(keypts_xyp)
-    
-    return list_kpts_per_crop
+
+    return list_kpts_per_crop

pre-requirements.txt

@@ -0,0 +1 @@
+pip>=26.2

pyproject.toml

@@ -0,0 +1,33 @@
+[project]
+name = "deeplabcut-modelzoo-superanimals"
+version = "0.1.0"
+description = "Gradio Space for the DeepLabCut Model Zoo SuperAnimal models"
+readme = "README.md"
+requires-python = ">=3.12,<3.13"
+classifiers = [
+  "Programming Language :: Python :: 3 :: Only",
+  "Programming Language :: Python :: 3.12",
+]
+# Hugging Face Spaces install from requirements.txt only: keep both lists in sync.
+dependencies = [
+  "deeplabcut[modelzoo,tf]==3.0.2",
+  "deeplabcut-live",
+  "dlclibrary",
+  "gitpython>=3.1.30",
+  "gradio==6.29",
+  "humanfriendly",
+  "psutil",
+  "ruamel-yaml==0.17.21",
+  "seaborn",
+  "ultralytics",
+]
+
+[tool.ruff]
+target-version = "py312"
+line-length = 120
+format.docstring-code-format = true
+lint.select = [ "B", "E", "F", "I", "PIE", "UP" ]
+lint.ignore = [ "B007", "E741" ]
+
+[tool.pyproject-fmt]
+max_supported_python = "3.12"

pytorch_utils.py

@@ -0,0 +1,118 @@
+import threading
+from pathlib import Path
+
+import numpy as np
+import PIL
+from deeplabcut.pose_estimation_pytorch.apis.utils import get_inference_runners
+from deeplabcut.pose_estimation_pytorch.config.pose import PoseConfig
+from deeplabcut.pose_estimation_pytorch.modelzoo.utils import MODEL_FILENAME_MAPPING
+from dlclibrary import download_huggingface_model
+
+# SuperAnimal (pose model, detector) used by the PyTorch backend
+PYTORCH_MODELS = {
+    "superanimal_quadruped": ("hrnet_w32", "fasterrcnn_resnet50_fpn_v2"),
+    "superanimal_topviewmouse": ("hrnet_w32", "fasterrcnn_resnet50_fpn_v2"),
+}
+
+MAX_INDIVIDUALS = 10
+MAX_IMAGE_SIZE = 1280  # longest side fed to the models (and drawn on)
+# next to the TF models, not deeplabcut/modelzoo/checkpoints: site-packages is read-only for the Space's non-root user
+WEIGHTS_DIR = Path(__file__).parent / "DLC_models" / "pytorch"
+
+_runners = {}
+_build_lock = threading.Lock()
+
+
+##########################################
+def snapshot_path(superanimal, model_name):
+    """Path to a SuperAnimal snapshot in WEIGHTS_DIR, downloaded on first use (as deeplabcut does)."""
+    name = f"{superanimal}_{model_name}"
+    path = WEIGHTS_DIR / f"{name}.pt"
+    if not path.exists():
+        source = MODEL_FILENAME_MAPPING.get(name, path.name)
+        rename = None if source == path.name else {source: path.name}
+        download_huggingface_model(name, target_dir=str(WEIGHTS_DIR), rename_mapping=rename)
+    return path
+
+
+##########################################
+def load_superanimal(superanimal, device="auto"):
+    """Build (once) the detector and pose runners for a SuperAnimal model; weights are downloaded on first use."""
+    with _build_lock:
+        if superanimal not in _runners:
+            pose_model, detector = PYTORCH_MODELS[superanimal]
+            cfg = PoseConfig.build_for_superanimal_inference(
+                superanimal,
+                model_name=pose_model,
+                detector_name=detector,
+                max_individuals=MAX_INDIVIDUALS,
+                device=device,
+            )
+            # keep low-score boxes: the UI threshold filters them afterwards
+            cfg["detector"]["model"]["box_score_thresh"] = 0.05
+            pose_runner, detector_runner = get_inference_runners(
+                cfg,
+                snapshot_path=snapshot_path(superanimal, pose_model),
+                detector_path=snapshot_path(superanimal, detector),
+                max_individuals=MAX_INDIVIDUALS,
+                inference_cfg={"multithreading": {"enabled": False}},
+            )
+            _runners[superanimal] = {
+                "pose": pose_runner,
+                "detector": detector_runner,
+                "bodyparts": list(cfg["metadata"]["bodyparts"]),
+                "lock": threading.Lock(),
+            }  # runners are not thread-safe
+    return _runners[superanimal]
+
+
+##########################################
+def resize_max_side(img, max_size=MAX_IMAGE_SIZE):
+    scale = max_size / max(img.size)
+    if scale >= 1:
+        return img
+    return img.resize([int(x * scale) for x in img.size], PIL.Image.Resampling.LANCZOS)
+
+
+##########################################
+def predict_superanimal(img_input, superanimal, bbox_likelihood_th, kpts_likelihood_th, full_image=False):
+    """Detect animals and estimate their pose with a PyTorch SuperAnimal model.
+
+    Returns the (resized) RGB image the predictions refer to, the list of animals
+    as dicts {'bbox': [x1,y1,x2,y2,conf], 'kpts': (num_keypoints, 3) array of x,y,llk
+    in image pixels, NaN below kpts_likelihood_th}, and the bodypart names.
+    """
+    img = resize_max_side(img_input.convert("RGB"))
+    img_np = np.asarray(img)
+    runners = load_superanimal(superanimal)
+
+    with runners["lock"]:
+        if full_image:
+            # skip the detector and treat the whole image as one animal
+            h, w = img_np.shape[:2]
+            detections = {
+                "bboxes": np.array([[0, 0, w, h]], dtype=np.float32),
+                "bbox_scores": np.array([1.0], dtype=np.float32),
+            }
+        else:
+            detections = runners["detector"].inference([img_np])[0]  # bboxes in xywh
+            keep = detections["bbox_scores"] >= bbox_likelihood_th
+            detections = {"bboxes": detections["bboxes"][keep], "bbox_scores": detections["bbox_scores"][keep]}
+
+        if len(detections["bboxes"]) == 0:
+            return img, [], runners["bodyparts"]
+
+        predictions = runners["pose"].inference([(img_np, detections)])[0]
+
+    animals = []
+    # outputs are padded to MAX_INDIVIDUALS with -1
+    for kpts, (x, y, w, h), score in zip(
+        predictions["bodyparts"], predictions["bboxes"], predictions["bbox_scores"], strict=True
+    ):
+        if score < 0:
+            continue
+        kpts = kpts.astype(float)
+        kpts[kpts[:, 2] < kpts_likelihood_th, :] = np.nan
+        animals.append({"bbox": [float(x), float(y), float(x + w), float(y + h), float(score)], "kpts": kpts})
+
+    return img, animals, runners["bodyparts"]

requirements.txt

@@ -1,10 +1,13 @@
-gradio
+# Hugging Face Spaces install from this file; keep in sync with pyproject.toml dependencies.
+# CPU-only PyTorch wheels: the Space has no GPU, and the default Linux wheels pull ~2-3 GB of CUDA libraries
+--extra-index-url https://download.pytorch.org/whl/cpu
+gradio==6.29.0
 gitpython>=3.1.30
 seaborn
-deeplabcut[modelzoo,tf]>=3.0.0rc14
+deeplabcut[modelzoo,tf]==3.0.2
 deeplabcut-live
 ruamel.yaml==0.17.21
 dlclibrary
 humanfriendly
 psutil
-ultralytics
+ultralytics

save_results.py

@@ -1,56 +0,0 @@
-import json 
-import numpy as np
-import pdb
-
-dict_pred = {0: 'animal', 1: 'person', 2: 'vehicle'}
-
-
-def save_results(md_results, dlc_outputs,map_label_id_to_str,thr,output_file = 'dowload_predictions.json'):
-
-    """
-
-    write json
-
-    """
-    info = {}
-    ## info megaDetector
-    info['file']= md_results.files[0]
-    number_bb = len(md_results.xyxy[0].tolist())
-    info['number_of_bb'] = number_bb
-    number_bb_thr = len(dlc_outputs)
-    labels = [n for n in map_label_id_to_str.values()]
-    #pdb.set_trace()
-    new_index = []
-    for i in range(number_bb):
-        corner_x1,corner_y1,corner_x2,corner_y2,confidence, _ =  md_results.xyxy[0].tolist()[i]
-
-        if confidence > thr:
-            new_index.append(i)
-
-
-    for i in range(number_bb_thr):
-        aux={}
-        corner_x1,corner_y1,corner_x2,corner_y2,confidence, _ =  md_results.xyxy[0].tolist()[new_index[i]]
-        aux['corner_1'] = (corner_x1,corner_y1)
-        aux['corner_2'] = (corner_x2,corner_y2)
-        aux['predict MD'] = md_results.names[0]
-        aux['confidence MD'] = confidence
-        
-        ## info dlc
-        kypts = []
-        for s in dlc_outputs[i]:
-            aux1 = []
-            for j in s:
-                aux1.append(float(j))
-
-            kypts.append(aux1)
-        aux['dlc_pred']  = dict(zip(labels,kypts))
-        info['bb_' + str(new_index[i]) ]=aux
-
-
-    with open(output_file, 'w') as f:
-        json.dump(info, f, indent=1)
-        print('Output file saved at {}'.format(output_file))
-
-    return output_file
-

ui_utils.py

@@ -1,21 +1,78 @@
 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(md_models_list, dlc_models_list):
+def gradio_inputs_for_MD_DLC(backends_list, md_models_list, dlc_models_list):
     # Input image
     gr_image_input = gr.Image(type="pil", label="Input Image")
 
     # Models
+    gr_backend_input = gr.Radio(
+        choices=backends_list,
+        value=backends_list[0],
+        label="Select backend",
+    )
+
     gr_mega_model_input = gr.Dropdown(
         choices=md_models_list,
         value="md_v5a",
         type="value",
-        label="Select Detector model",
+        label="Select Detector model (TensorFlow legacy only)",
+        visible=gr_backend_input.value != "PyTorch",
     )
 
     gr_dlc_model_input = gr.Dropdown(
         choices=dlc_models_list,
-        value="superanimal_quadruped_dlcrnet",
+        value="superanimal_quadruped",
         type="value",
         label="Select DeepLabCut model",
     )
@@ -23,12 +80,7 @@ def gradio_inputs_for_MD_DLC(md_models_list, dlc_models_list):
     # Other inputs
     gr_dlc_only_checkbox = gr.Checkbox(
         value=False,
-        label="Run DLClive only, directly on input image?",
-    )
-
-    gr_str_labels_checkbox = gr.Checkbox(
-        value=True,
-        label="Show bodypart labels?",
+        label="Run DeepLabCut only, directly on input image?",
     )
 
     # Gradio Slider signature is (minimum, maximum, value, step, ...)
@@ -49,36 +101,60 @@ def gradio_inputs_for_MD_DLC(md_models_list, dlc_models_list):
     )
 
     # Data viz
-    gr_keypt_color = gr.ColorPicker(
-        value="#862db7",
-        label="Choose color for keypoint label",
-    )
-
-    gr_labels_font_style = gr.Dropdown(
-        choices=["amiko", "animals", "nature", "painter", "zen"],
-        value="amiko",
-        type="value",
-        label="Select keypoint label font",
-    )
-
-    gr_slider_font_size = gr.Slider(
-        minimum=5,
-        maximum=30,
-        value=8,
-        step=1,
-        label="Set font size",
-    )
-
-    gr_slider_marker_size = gr.Slider(
-        minimum=1,
-        maximum=20,
-        value=9,
-        step=1,
-        label="Set marker size",
-    )
+    with gr.Accordion("Display options", open=False):
+        gr_str_labels_checkbox = gr.Checkbox(
+            value=True,
+            label="Show bodypart labels?",
+        )
+
+        gr_color_by_confidence_checkbox = gr.Checkbox(
+            value=True,
+            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",
+            type="value",
+            label="Select keypoint label font",
+        )
+
+        gr_slider_font_size = gr.Slider(
+            minimum=5,
+            maximum=30,
+            value=18,
+            step=1,
+            label="Set font size",
+        )
+
+        gr_slider_marker_size = gr.Slider(
+            minimum=1,
+            maximum=20,
+            value=6,
+            step=1,
+            label="Set marker size",
+        )
 
     return [
         gr_image_input,
+        gr_backend_input,
         gr_mega_model_input,
         gr_dlc_model_input,
         gr_dlc_only_checkbox,
@@ -89,38 +165,64 @@ def gradio_inputs_for_MD_DLC(md_models_list, dlc_models_list):
         gr_slider_font_size,
         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_file_download = gr.File(label="Download JSON file")
-    return [gr_image_output, gr_file_download]
+    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")
+    gr_confidence_table = gr.Dataframe(
+        headers=["animal", "bodypart", "confidence"],
+        label="Keypoint confidence (lowest first)",
+        interactive=False,
+    )
+    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():
-    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 = [[
-        "examples/dog.jpeg",
-        "md_v5a",
-        "superanimal_quadruped_dlcrnet",
-        False,
-        True,
-        0.5,
-        0.0,
-        "amiko",
-        9,
-        "#ff0000",
-        3,
-    ]]
-
-    return [title, description, examples]
+    examples = [
+        [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",
+        )
+        for font_size, marker_size in [example_sizes(image)]
+    ]
+
+    return [title, description, examples]

viz_utils.py

@@ -1,31 +1,38 @@
-import json 
-import numpy as np
+import json
+import math
+from datetime import date
 
-from matplotlib import cm
-import matplotlib
-from PIL import Image, ImageColor, ImageFont, ImageDraw 
 import numpy as np
-import pdb
-from datetime import date
+from matplotlib import colormaps
+from PIL import ImageColor, ImageDraw, ImageFont
+
 today = date.today()
-FONTS = {'amiko': "fonts/Amiko-Regular.ttf",
-        'nature': "fonts/LoveNature.otf", 
-        'painter':"fonts/PainterDecorator.otf",
-        'animals': "fonts/UncialAnimals.ttf", 
-        'zen': "fonts/ZEN.TTF"}
+FONTS = {
+    "amiko": "fonts/Amiko-Regular.ttf",
+    "nature": "fonts/LoveNature.otf",
+    "painter": "fonts/PainterDecorator.otf",
+    "animals": "fonts/UncialAnimals.ttf",
+    "zen": "fonts/ZEN.TTF",
+}
+# perceptually uniform maps first; turbo separates neighbouring bodyparts best
+COLORMAPS = ["viridis", "plasma", "magma", "cividis", "turbo"]
+
 
 #########################################
 #  Draw keypoints on image
-def draw_keypoints_on_image(image,
-                            keypoints,
-                            map_label_id_to_str,
-                            flag_show_str_labels,
-                            use_normalized_coordinates=True,                           
-                            font_style='amiko',
-                            font_size=8,
-                            keypt_color="#ff0000",
-                            marker_size=2,
-                            ):
+def draw_keypoints_on_image(
+    image,
+    keypoints,
+    map_label_id_to_str,
+    flag_show_str_labels,
+    use_normalized_coordinates=True,
+    font_style="amiko",
+    font_size=8,
+    keypt_color="#ff0000",
+    marker_size=2,
+    color_by_confidence=True,
+    colormap="viridis",
+):
     """Draws keypoints on an image.
     Modified from:
         https://www.programcreek.com/python/?code=fjchange%2Fobject_centric_VAD%2Fobject_centric_VAD-master%2Fobject_detection%2Futils%2Fvisualization_utils.py
@@ -39,159 +46,232 @@ def draw_keypoints_on_image(image,
     use_normalized_coordinates: if True (default), treat keypoint values as
         relative to the image.  Otherwise treat them as absolute.
 
-    
+
     """
     # get a drawing context
-    draw = ImageDraw.Draw(image,"RGBA")  
+    draw = ImageDraw.Draw(image, "RGBA")
 
     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:
         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')
-    cmap2 = matplotlib.cm.get_cmap('Greys')
+
+    cmap = colormaps[colormap]
     # 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!
+    for i, (keypoint_x, keypoint_y) in enumerate(zip(keypoints_x, keypoints_y, strict=True)):
         # 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]]
-        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]
+
+        confidence = float(np.clip(confidences[i], 0, 1))
+        if color_by_confidence:
+            # 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(cmap(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]
 
         # add string labels around keypoints
         if flag_show_str_labels:
-            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],
-                      ImageColor.getcolor(keypt_color, "RGB"), # rgb #
-                      font=font)
+            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), #-------
+                display_bodypart(map_label_id_to_str[i]),
+                ImageColor.getcolor(keypt_color, "RGB"),  # rgb #
+                font=font,
+            )
+
+
+#########################################
+#  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",
+}
+
+
+def display_bodypart(name):
+    return BODYPART_DISPLAY_NAMES.get(name, name.replace("_", " "))
+
+
+#########################################
+#  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, display_bodypart(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,
-                     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")
-
-    return imgR
+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)))
+
+    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"
+
 
 ###########################################
-def save_results_as_json(md_results, dlc_outputs, map_dlc_label_id_to_str, thr,model,mega_model_input, path_to_output_file = 'download_predictions.json'):
+#  JSON outputs: pixel coordinates in the input image, hidden keypoints as null
+COORDINATES = "pixels in the input image (after EXIF orientation), origin top-left, x right, y down"
 
-    """
-    Output detections as json file
 
-    """
-    # initialise dict to save to json
-    info = {}
-    info['date'] = str(today)
-    info['MD_model'] = str(mega_model_input)
-    # info from megaDetector
-    info['file']= md_results.files[0]
-    number_bb = len(md_results.xyxy[0].tolist())
-    info['number_of_bb'] = number_bb
-    # info from DLC
-    number_bb_thr = len(dlc_outputs)
-    labels = [n for n in map_dlc_label_id_to_str.values()]
-    
-    # create list of bboxes above th
-    new_index = []
-    for i in range(number_bb):
-        corner_x1,corner_y1,corner_x2,corner_y2,confidence, _ =  md_results.xyxy[0].tolist()[i]
-
-        if confidence > thr:
-            new_index.append(i)
-
-    # define aux dict for every bounding box above threshold
-    for i in range(number_bb_thr):
-        aux={}
-        # MD output
-        corner_x1,corner_y1,corner_x2,corner_y2,confidence, _ =  md_results.xyxy[0].tolist()[new_index[i]]
-        aux['corner_1'] = (corner_x1,corner_y1)
-        aux['corner_2'] = (corner_x2,corner_y2)
-        aux['predict MD'] = md_results.names[0]
-        aux['confidence MD'] = confidence
-        
-        # DLC output
-        info['dlc_model'] = model
-        kypts = []
-        for s in dlc_outputs[i]:
-            aux1 = []
-            for j in s:
-                aux1.append(float(j))
-
-            kypts.append(aux1)
-        aux['dlc_pred']  = dict(zip(labels,kypts))
-        info['bb_' + str(new_index[i]) ]=aux
-
-    # save dict as json
-    with open(path_to_output_file, 'w') as f:
-        json.dump(info, f, indent=1)
-        print('Output file saved at {}'.format(path_to_output_file))
+def keypoints_to_json(kpts, offset=(0, 0), scale=(1.0, 1.0)):
+    """[x, y, confidence] per keypoint, mapped by (k + offset) * scale; NaN (below threshold) becomes null."""
+    out = []
+    for x, y, conf in np.asarray(kpts, dtype=float)[:, :3]:
+        if np.isnan(conf):
+            out.append(None)
+        else:
+            out.append([(x + offset[0]) * scale[0], (y + offset[1]) * scale[1], conf])
+    return out
+
 
+def write_json(info, path_to_output_file):
+    with open(path_to_output_file, "w") as f:
+        json.dump(info, f, indent=1, allow_nan=False)
     return path_to_output_file
 
 
-def save_results_only_dlc(dlc_outputs,map_label_id_to_str,model,output_file = 'dowload_predictions_dlc.json'):
+def json_header(image_size, annotated_size):
+    return {
+        "date": str(today),
+        "coordinates": COORDINATES,
+        "image_size": list(image_size),
+        "annotated_image_size": list(annotated_size),
+    }
 
+
+def save_results_as_json(
+    md_results,
+    dlc_outputs,
+    animal_bboxes,
+    map_dlc_label_id_to_str,
+    model,
+    mega_model_input,
+    image_size,
+    path_to_output_file="download_predictions.json",
+):
+    """TF (legacy) MegaDetector + DLC results.
+
+    animal_bboxes: detection rows [x1,y1,x2,y2,conf,label] in the MegaDetector frame, one per entry of dlc_outputs
+    dlc_outputs: keypoints relative to each crop (crops start at floor(x1), floor(y1), see crop_animal_detections)
     """
-    write json dlc output
-    """
-    info = {}
-    info['date'] = str(today)
-    labels = [n for n in map_label_id_to_str.values()]
-    info['dlc_model'] = model
-    kypts = []
-    for s in dlc_outputs:
-        aux1 = []
-        for j in s:
-            aux1.append(float(j))
+    md_h, md_w = md_results.ims[0].shape[:2]
+    scale = (image_size[0] / md_w, image_size[1] / md_h)
+    info = json_header(image_size, (md_w, md_h))
+    info["MD_model"] = str(mega_model_input)
+    info["number_of_bb"] = len(dlc_outputs)
+    info["dlc_model"] = model
+    labels = list(map_dlc_label_id_to_str.values())
 
-        kypts.append(aux1)
-    info['dlc_pred']  = dict(zip(labels,kypts))
+    for i, kpts in enumerate(dlc_outputs):
+        x1, y1, x2, y2, confidence, _ = animal_bboxes[i]
+        info["bb_" + str(i)] = {
+            "corner_1": (x1 * scale[0], y1 * scale[1]),
+            "corner_2": (x2 * scale[0], y2 * scale[1]),
+            "predict MD": md_results.names[0],
+            "confidence MD": float(confidence),
+            "dlc_pred": dict(
+                zip(labels, keypoints_to_json(kpts, offset=(math.floor(x1), math.floor(y1)), scale=scale), strict=True)
+            ),
+        }
+    return write_json(info, path_to_output_file)
 
-    with open(output_file, 'w') as f:
-        json.dump(info, f, indent=1)
-        print('Output file saved at {}'.format(output_file))
 
-    return output_file
+def save_results_only_dlc(
+    dlc_outputs, map_label_id_to_str, model, image_size, output_file="dowload_predictions_dlc.json"
+):
+    """TF (legacy) DLC run on the whole input image (keypoints already in input pixels)."""
+    info = json_header(image_size, image_size)
+    info["dlc_model"] = model
+    info["dlc_pred"] = dict(zip(map_label_id_to_str.values(), keypoints_to_json(dlc_outputs), strict=True))
+    return write_json(info, output_file)
 
 
-###########################################
+def save_results_pytorch(
+    animals,
+    map_label_id_to_str,
+    model,
+    pose_model,
+    detector,
+    image_size,
+    annotated_size,
+    path_to_output_file="download_predictions.json",
+):
+    """PyTorch SuperAnimal results (same layout as save_results_as_json).
+
+    animals: list of {'bbox': [x1,y1,x2,y2,conf], 'kpts': (num_keypoints, 3)}, in the annotated (resized) image
+    detector: None if the detector was skipped (whole image used as one animal)
+    """
+    scale = (image_size[0] / annotated_size[0], image_size[1] / annotated_size[1])
+    info = json_header(image_size, annotated_size)
+    info["backend"] = "pytorch"
+    info["dlc_model"] = model
+    info["pose_model"] = pose_model
+    info["detector"] = detector
+    info["number_of_bb"] = len(animals)
+    labels = list(map_label_id_to_str.values())
+
+    for i, animal in enumerate(animals):
+        x1, y1, x2, y2, confidence = animal["bbox"]
+        info["bb_" + str(i)] = {
+            "corner_1": (x1 * scale[0], y1 * scale[1]),
+            "corner_2": (x2 * scale[0], y2 * scale[1]),
+            "confidence": confidence,
+            "dlc_pred": dict(zip(labels, keypoints_to_json(animal["kpts"], scale=scale), strict=True)),
+        }
+    return write_json(info, path_to_output_file)
+
+
+###########################################