run pre commit hooks

README.md

@@ -10,4 +10,4 @@ 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,51 +1,55 @@
-# 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 threading
-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
-
-from viz_utils import save_results_as_json, draw_keypoints_on_image, draw_bbox_w_text, save_results_only_dlc, save_results_pytorch
-from viz_utils import add_confidence_legend, keypoint_confidence_rows, save_annotated_image
-from detection_utils import predict_md, crop_animal_detections
-from dlc_utils import predict_dlc
-from pytorch_utils import predict_superanimal, load_superanimal, PYTORCH_MODELS
-from ui_utils import gradio_inputs_for_MD_DLC, gradio_outputs_for_MD_DLC, gradio_description_and_examples
 
-from deeplabcut.utils import auxiliaryfunctions
+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
 
+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 viz_utils import (
+    add_confidence_legend,
+    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,
+)
 
 # TESTING (passes) download the SuperAnimal models:
-#model = 'superanimal_topviewmouse'
-#train_dir = 'DLC_models/sa-tvm'
-#download_huggingface_model(model, train_dir)
+# 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",
+}
 
 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")}
-
+DLC_models_dict = {
+    "superanimal_topviewmouse": ("superanimal_topviewmouse_dlcrnet", "DLC_models/sa-tvm"),
+    "superanimal_quadruped": ("superanimal_quadruped_dlcrnet", "DLC_models/sa-q"),
+}
 
 
 #####################################################
@@ -59,99 +63,102 @@ def finalize_outputs(img_output, download_file, kpts_per_animal, map_label_id_to
 
 
 #####################################################
-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,
-                             ):
+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,
+):
     # 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)
+    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)
+        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,
+        )
         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)
 
     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)
+    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
+    )
 
 
 #####################################################
-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,
-                     ):
+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,
+):
 
     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)
+        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,
+        )
 
     # 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 (and their bboxes) with confidence above th
-        list_crops, list_bboxes = crop_animal_detections(img_input,
-                                                         md_results,
-                                                         bbox_likelihood_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?
     path_to_DLCmodel = dlc_model_dir
@@ -159,124 +166,114 @@ def predict_pipeline(img_input,
         download_huggingface_model(dlc_model_name, path_to_DLCmodel)
 
     # extract map label ids to strings
-    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) 
-    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,
-                                color_by_confidence=flag_color_by_confidence)
-
-        donw_file = save_results_only_dlc(list_kpts_per_crop[0], map_label_id_to_str,dlc_model_name)
-
-        return finalize_outputs(img_input, donw_file,
-                                [list_kpts_per_crop[0]], map_label_id_to_str,
-                                flag_color_by_confidence)
+        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,
+        )
+
+        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
+        )
 
     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 np_crop, kpts_crop, bb_per_animal in zip(list_crops,
-                                                     list_kpts_per_crop,
-                                                     list_bboxes):
+        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,
-                                    color_by_confidence=flag_color_by_confidence)
+            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,
+            )
 
             # Paste crop in original image
-            img_background.paste(img_crop,
-                                 box = tuple([int(t) for t in bb_per_animal[:2]]))
+            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)  # TODO: add selectable color for bbox?
 
         # Save detection results as json
-        download_file  = save_results_as_json(md_results,list_kpts_per_crop,list_bboxes,map_label_id_to_str,dlc_model_name,mega_model_input)
-
-        return finalize_outputs(img_background, download_file,
-                                list_kpts_per_crop, map_label_id_to_str,
-                                flag_color_by_confidence)
+        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 finalize_outputs(
+            img_background, download_file, list_kpts_per_crop, map_label_id_to_str, flag_color_by_confidence
+        )
 
 
 #########################################################
 # Define user interface and launch
-[gr_title,
- gr_description,
- examples] = gradio_description_and_examples()
+[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()))
+            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)
+    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")
+    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

detection_utils.py

@@ -1,86 +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
+    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
-                              )
+    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_np_animal_crops = []
-    list_animal_bboxes = [] # detection rows [x1,y1,x2,y2,conf,label] matching each crop
+    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())
+                list_animal_bboxes.append(det_array[j, :].tolist())
 
-    return list_np_animal_crops, list_animal_bboxes
+    return list_np_animal_crops, list_animal_bboxes

dlc_utils.py

@@ -1,16 +1,10 @@
-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 []
@@ -25,10 +19,10 @@ def predict_dlc(list_np_crops,
 
     list_kpts_per_crop = []
     for crop in list_np_crops:
-        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
-        keypts_xyp[keypts_xyp[:,-1] < kpts_likelihood_th,:] = np.nan
+        keypts_xyp[keypts_xyp[:, -1] < kpts_likelihood_th, :] = np.nan
         # add kpts of this crop to list
         list_kpts_per_crop.append(keypts_xyp)
 
-    return list_kpts_per_crop
+    return list_kpts_per_crop

pytorch_utils.py

@@ -6,10 +6,11 @@ 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')}
+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)
@@ -27,11 +28,13 @@ def load_superanimal(superanimal, device="auto"):
     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)
+            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(
@@ -41,10 +44,12 @@ def load_superanimal(superanimal, device="auto"):
                 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
+            _runners[superanimal] = {
+                "pose": pose_runner,
+                "detector": detector_runner,
+                "bodyparts": list(cfg["metadata"]["bodyparts"]),
+                "lock": threading.Lock(),
+            }  # runners are not thread-safe
     return _runners[superanimal]
 
 
@@ -57,11 +62,7 @@ def resize_max_side(img, max_size=MAX_IMAGE_SIZE):
 
 
 ##########################################
-def predict_superanimal(img_input,
-                        superanimal,
-                        bbox_likelihood_th,
-                        kpts_likelihood_th,
-                        full_image=False):
+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
@@ -76,13 +77,14 @@ def predict_superanimal(img_input,
         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)}
+            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]}
+            detections = {"bboxes": detections["bboxes"][keep], "bbox_scores": detections["bbox_scores"][keep]}
 
         if len(detections["bboxes"]) == 0:
             return img, [], runners["bodyparts"]
@@ -91,14 +93,13 @@ def predict_superanimal(img_input,
 
     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"]):
+    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})
+        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,3 +1,4 @@
+# Hugging Face Spaces install from this file; keep in sync with pyproject.toml dependencies.
 gradio
 gitpython>=3.1.30
 seaborn
@@ -7,4 +8,4 @@ ruamel.yaml==0.17.21
 dlclibrary
 humanfriendly
 psutil
-ultralytics
+ultralytics

ui_utils.py

@@ -134,4 +134,4 @@ def gradio_description_and_examples():
         for image in ("examples/dog.jpeg", "examples/cat.jpg")
     ]
 
-    return [title, description, examples]
+    return [title, description, examples]

viz_utils.py

@@ -1,32 +1,34 @@
-import json 
-import numpy as np
+import json
+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 cm
+from PIL import Image, 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",
+}
+
 
 #########################################
 #  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,
-                            color_by_confidence=True,
-                            ):
+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,
+):
     """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
@@ -40,10 +42,10 @@ 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]
@@ -54,10 +56,10 @@ def draw_keypoints_on_image(image,
     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')
+
+    # cmap = matplotlib.cm.get_cmap('hsv')
     # draw ellipses around keypoints
-    for i, (keypoint_x, keypoint_y) in enumerate(zip(keypoints_x, keypoints_y)):
+    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
@@ -71,22 +73,30 @@ def draw_keypoints_on_image(image,
             round_fill = list(cm.viridis(i / max(len(keypoints) - 1, 1), bytes=True))
             round_fill[3] = round(confidence * 255)
             round_fill = tuple(round_fill)
-        draw.ellipse([(keypoint_x - marker_size, keypoint_y - marker_size),
-                      (keypoint_x + marker_size, keypoint_y + marker_size)],
-                       fill=tuple(round_fill), outline= 'black', width=1) #fill and outline: [0,255]
+        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), #-------
+                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'):
+def add_confidence_legend(image, font_style="amiko"):
     """Returns the image with a white band below it holding a viridis strip (confidence 0 to 1).
 
     The band is added below the image, so the legend never covers it and keypoint
@@ -105,8 +115,7 @@ def add_confidence_legend(image, font_style='amiko'):
     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))
+        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")
@@ -128,39 +137,44 @@ def keypoint_confidence_rows(kpts_per_animal, map_label_id_to_str):
 
 #########################################
 #  Save the annotated image for download
-def save_annotated_image(image, path_to_output_file='download_annotated.png'):
+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()
+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
+    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) 
+    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 )
+    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")
+    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 save_results_as_json(md_results, dlc_outputs, animal_bboxes, map_dlc_label_id_to_str, model,mega_model_input, path_to_output_file = 'download_predictions.json'):
 
+###########################################
+def save_results_as_json(
+    md_results,
+    dlc_outputs,
+    animal_bboxes,
+    map_dlc_label_id_to_str,
+    model,
+    mega_model_input,
+    path_to_output_file="download_predictions.json",
+):
     """
     Output detections as json file
 
@@ -168,26 +182,26 @@ def save_results_as_json(md_results, dlc_outputs, animal_bboxes, map_dlc_label_i
     """
     # initialise dict to save to json
     info = {}
-    info['date'] = str(today)
-    info['MD_model'] = str(mega_model_input)
+    info["date"] = str(today)
+    info["MD_model"] = str(mega_model_input)
     # info from megaDetector
-    info['file']= md_results.files[0]
+    info["file"] = md_results.files[0]
     number_bb = len(md_results.xyxy[0].tolist())
-    info['number_of_bb'] = number_bb
+    info["number_of_bb"] = number_bb
     # info from DLC
-    info['dlc_model'] = model
+    info["dlc_model"] = model
     labels = [n for n in map_dlc_label_id_to_str.values()]
 
     # define aux dict for every animal bounding box above threshold
     for i in range(len(dlc_outputs)):
-        aux={}
+        aux = {}
         # MD output
-        corner_x1,corner_y1,corner_x2,corner_y2,confidence, _ =  animal_bboxes[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
-        
+        corner_x1, corner_y1, corner_x2, corner_y2, confidence, _ = animal_bboxes[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
         kypts = []
         for s in dlc_outputs[i]:
@@ -196,26 +210,25 @@ def save_results_as_json(md_results, dlc_outputs, animal_bboxes, map_dlc_label_i
                 aux1.append(float(j))
 
             kypts.append(aux1)
-        aux['dlc_pred']  = dict(zip(labels,kypts))
-        info['bb_' + str(i) ]=aux
+        aux["dlc_pred"] = dict(zip(labels, kypts, strict=True))
+        info["bb_" + str(i)] = aux
 
     # save dict as json
-    with open(path_to_output_file, 'w') as f:
+    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))
+        print(f"Output file saved at {path_to_output_file}")
 
     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 save_results_only_dlc(dlc_outputs, map_label_id_to_str, model, output_file="dowload_predictions_dlc.json"):
     """
     write json dlc output
     """
     info = {}
-    info['date'] = str(today)
+    info["date"] = str(today)
     labels = [n for n in map_label_id_to_str.values()]
-    info['dlc_model'] = model
+    info["dlc_model"] = model
     kypts = []
     for s in dlc_outputs:
         aux1 = []
@@ -223,17 +236,18 @@ def save_results_only_dlc(dlc_outputs,map_label_id_to_str,model,output_file = 'd
             aux1.append(float(j))
 
         kypts.append(aux1)
-    info['dlc_pred']  = dict(zip(labels,kypts))
+    info["dlc_pred"] = dict(zip(labels, kypts, strict=True))
 
-    with open(output_file, 'w') as f:
+    with open(output_file, "w") as f:
         json.dump(info, f, indent=1)
-        print('Output file saved at {}'.format(output_file))
+        print(f"Output file saved at {output_file}")
 
     return output_file
 
 
-def save_results_pytorch(animals, map_label_id_to_str, model, pose_model, detector, path_to_output_file = 'download_predictions.json'):
-
+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)
 
@@ -241,28 +255,28 @@ def save_results_pytorch(animals, map_label_id_to_str, model, pose_model, detect
     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)
+    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']
+        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
+        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"]], strict=True))
+        info["bb_" + str(i)] = aux
 
-    with open(path_to_output_file, 'w') as f:
+    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))
+        print(f"Output file saved at {path_to_output_file}")
 
     return path_to_output_file
 
 
-###########################################
+###########################################