Add PyTorch SuperAnimal backend

Introduce a PyTorch inference path for SuperAnimal models, including backend selection in the UI, cached pose/detector runners, and JSON export of predictions. The legacy TensorFlow flow now maps model names separately from local directories, and MegaDetector device selection was corrected for yolov5 loading. The DeepLabCut dependency is also pinned to 3.0.2 to match the new backend integration.

app.py

@@ -16,9 +16,10 @@ import dlclive
 from PIL import Image, ImageColor, ImageFont, ImageDraw 
 import requests
 
-from viz_utils import save_results_as_json, draw_keypoints_on_image, draw_bbox_w_text, save_results_only_dlc
+from viz_utils import save_results_as_json, draw_keypoints_on_image, draw_bbox_w_text, save_results_only_dlc, save_results_pytorch
 from detection_utils import predict_md, crop_animal_detections
 from dlc_utils import predict_dlc
+from pytorch_utils import predict_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
@@ -42,14 +43,58 @@ image = Image.open(requests.get(url, stream=True).raw)
 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_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,
+                             ):
+    # 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)
+        if not flag_dlc_only:
+            draw_bbox_w_text(img_output,
+                             animal['bbox'],
+                             font_size=font_size)
+
+    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
 
 
 #####################################################
 def predict_pipeline(img_input,
+                     backend,
                      mega_model_input,
                      dlc_model_input_str,
                      flag_dlc_only,
@@ -62,6 +107,21 @@ def predict_pipeline(img_input,
                      marker_size,
                      ):
 
+    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)
+
+    # 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
@@ -81,15 +141,12 @@ def predict_pipeline(img_input,
     
     # 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_path = os.path.join(dlc_model_dir,
                                  'pose_cfg.yaml')
     with open(pose_cfg_path, "r") as stream:
         pose_cfg_dict = yaml.safe_load(stream) 
@@ -119,7 +176,7 @@ def predict_pipeline(img_input,
                                 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)
+        donw_file = save_results_only_dlc(list_kpts_per_crop[0], map_label_id_to_str,dlc_model_name)
 
         return img_input, donw_file
 
@@ -163,7 +220,7 @@ 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_input_str,mega_model_input)
+        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
 
@@ -171,7 +228,8 @@ def predict_pipeline(img_input,
 
 #########################################################
 # Define user interface and launch
-inputs = gradio_inputs_for_MD_DLC(list(MD_models_dict.keys()),
+inputs = gradio_inputs_for_MD_DLC(BACKENDS,
+                                  list(MD_models_dict.keys()),
                                   list(DLC_models_dict.keys()))
 outputs = gradio_outputs_for_MD_DLC()                                    
 [gr_title, 

detection_utils.py

@@ -24,11 +24,8 @@ def predict_md(im,
     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')
+    # 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
@@ -37,12 +34,7 @@ def predict_md(im,
                               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)    
+                              )
 
     ## detect objects
     results = MD_model(im)  # inference # vars(results).keys()= dict_keys(['imgs', 'pred', 'names', 'files', 'times', 'xyxy', 'xywh', 'xyxyn', 'xywhn', 'n', 't', 's'])

pytorch_utils.py

@@ -0,0 +1,104 @@
+import threading
+
+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 get_super_animal_snapshot_path
+
+
+# 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)
+
+_runners = {}
+_build_lock = threading.Lock()
+
+
+##########################################
+def load_superanimal(superanimal, device="auto"):
+    """Build (once) the detector and pose runners for a SuperAnimal model.
+
+    Weights are downloaded on first use to deeplabcut/modelzoo/checkpoints.
+    """
+    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=get_super_animal_snapshot_path(superanimal, pose_model),
+                detector_path=get_super_animal_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"]):
+        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,7 +1,7 @@
 gradio
 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

ui_utils.py

@@ -1,21 +1,27 @@
 import gradio as gr
 
 
-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)",
     )
 
     gr_dlc_model_input = gr.Dropdown(
         choices=dlc_models_list,
-        value="superanimal_quadruped_dlcrnet",
+        value="superanimal_quadruped",
         type="value",
         label="Select DeepLabCut model",
     )
@@ -23,7 +29,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?",
+        label="Run DeepLabCut only, directly on input image?",
     )
 
     gr_str_labels_checkbox = gr.Checkbox(
@@ -79,6 +85,7 @@ def gradio_inputs_for_MD_DLC(md_models_list, dlc_models_list):
 
     return [
         gr_image_input,
+        gr_backend_input,
         gr_mega_model_input,
         gr_dlc_model_input,
         gr_dlc_only_checkbox,
@@ -111,8 +118,9 @@ def gradio_description_and_examples():
 
     examples = [[
         "examples/dog.jpeg",
+        "PyTorch",
         "md_v5a",
-        "superanimal_quadruped_dlcrnet",
+        "superanimal_quadruped",
         False,
         True,
         0.5,

viz_utils.py

@@ -185,4 +185,37 @@ def save_results_only_dlc(dlc_outputs,map_label_id_to_str,model,output_file = 'd
     return output_file
 
 
+def save_results_pytorch(animals, map_label_id_to_str, model, pose_model, detector, path_to_output_file = 'download_predictions.json'):
+
+    """
+    Output PyTorch SuperAnimal predictions as json file (same layout as save_results_as_json)
+
+    animals: list of {'bbox': [x1,y1,x2,y2,conf], 'kpts': (num_keypoints, 3)}, in image coords
+    detector: None if the detector was skipped (whole image used as one animal)
+    """
+    info = {}
+    info['date'] = str(today)
+    info['backend'] = 'pytorch'
+    info['dlc_model'] = model
+    info['pose_model'] = pose_model
+    info['detector'] = detector
+    info['number_of_bb'] = len(animals)
+    labels = [n for n in map_label_id_to_str.values()]
+
+    for i, animal in enumerate(animals):
+        corner_x1, corner_y1, corner_x2, corner_y2, confidence = animal['bbox']
+        aux = {}
+        aux['corner_1'] = (corner_x1, corner_y1)
+        aux['corner_2'] = (corner_x2, corner_y2)
+        aux['confidence'] = confidence
+        aux['dlc_pred'] = dict(zip(labels, [[float(v) for v in kpt] for kpt in animal['kpts']]))
+        info['bb_' + str(i)] = aux
+
+    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))
+
+    return path_to_output_file
+
+
 ###########################################