Update app requirements and code

.gitignore

@@ -1,16 +1,3 @@
 # Byte-compiled / optimized / DLL files
 __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/*/
+model/__pycache__/

.pre-commit-config.yaml

@@ -1,28 +0,0 @@
-# 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

@@ -4,10 +4,9 @@ emoji: 🐕🐁🐴😻🐘🐆🐿🐂🦘🦒
 colorFrom: blue
 colorTo: purple
 sdk: gradio
-python_version: 3.12
-sdk_version: 6.29.0
+sdk_version: 6.11.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,314 +1,192 @@
-# 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 gradio as gr
-import numpy as np
 import yaml
-from dlclibrary.dlcmodelzoo.modelzoo_download import (
-    download_huggingface_model,
-)
-from dlclive import Processor
+import numpy as np
+from matplotlib import cm
+import gradio as gr
+import deeplabcut
+import dlclibrary
+import dlclive
+import transformers
 
-# import transformers
-from PIL import Image
+from PIL import Image, ImageColor, ImageFont, ImageDraw 
+import requests
 
-from detection_utils import crop_animal_detections, predict_md
+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 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,
-)
-
-# 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",
-}
+from ui_utils import gradio_inputs_for_MD_DLC, gradio_outputs_for_MD_DLC, gradio_description_and_examples
 
-BACKENDS = ["PyTorch", "TensorFlow (legacy)"]
+from deeplabcut.utils import auxiliaryfunctions
+from dlclibrary.dlcmodelzoo.modelzoo_download import (
+    download_huggingface_model,
+    MODELOPTIONS,
+)
+from dlclive import DLCLive, Processor
 
-# 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"),
-}
 
+# TESTING (passes) download the SuperAnimal models:
+#model = 'superanimal_topviewmouse'
+#train_dir = 'DLC_models/sa-tvm'
+#download_huggingface_model(model, train_dir)
 
-#####################################################
-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
+# grab demo data cooco cat:
+url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+image = Image.open(requests.get(url, stream=True).raw)
 
+# 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"}
 
-#####################################################
-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,
-    )
+# DLC models target  dirs
+DLC_models_dict = {'superanimal_topviewmouse_dlcrnet': "DLC_models/sa-tvm",
+                   'superanimal_quadruped_dlcrnet': "DLC_models/sa-q"}
+                    
 
 
 #####################################################
-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]
+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,
+                     ):
 
     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)
+        # Obtain animal crops for bboxes with confidence above th
+        list_crops = 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
-    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)
+    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)
 
     # extract map label ids to strings
-    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,
-            )
-        ]
-    )
-
-    ##############################################################
+    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'])])
+
+
+##############################################################
     # 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,
-            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
-        )
+        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
 
     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]))
-
+        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):
+        for ic, (np_crop, kpts_crop) in enumerate(zip(list_crops,
+                                                      list_kpts_per_crop)):
+
             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,
-                colormap=colormap,
-            )
+            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)
 
             # 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 md_results.xyxy[0][ic,:2]]))
 
             # Plot bbox
-            draw_bbox_w_text(img_background, bb_per_animal, font_size=font_size, bbox_color=bbox_color)
-
+            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,
-            list_bboxes,
-            map_label_id_to_str,
-            dlc_model_name,
-            mega_model_input,
-            image_size=img_input.size,
-        )
+        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)         
+
+        return img_background, download_file
 
-        return finalize_outputs(
-            img_background, download_file, list_kpts_per_crop, map_label_id_to_str, flag_color_by_confidence, colormap
-        )
 
 
 #########################################################
 # Define user interface and launch
-[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())
+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,
+                    theme="huggingface")
+
+demo.launch(enable_queue=True, share=True)

detection_utils.py

@@ -1,74 +1,90 @@
-import math
 
-import numpy as np
-import PIL
+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 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: 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,
-    )
+    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
+                              force_reload=True,
+                              device=md_device) 
+                              
+    # send model to gpu if possible
+    if (md_device == torch.device('cuda')):
+        print('Sending model to GPU')
+        MD_model.to(md_device)    
 
     ## detect objects
-    # vars(results).keys(): imgs, pred, names, files, times, xyxy, xywh, xyxyn, xywhn, n, t, s
-    results = MD_model(im)
-
-    return results
+    results = MD_model(im)  # inference # vars(results).keys()= dict_keys(['imgs', 'pred', 'names', 'files', 'times', 'xyxy', 'xywh', 'xyxyn', 'xywhn', 'n', 't', 's'])
+    
+    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, list_animal_bboxes
+    return list_np_animal_crops

dlc_utils.py

@@ -1,28 +1,32 @@
+import deeplabcut
+from tkinter import W
+import gradio as gr
 import numpy as np
-from dlclive import DLCLive
+from dlclive import DLCLive, Processor
 
 
 ##########################################
-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]
-
+def predict_dlc(list_np_crops,
+                kpts_likelihood_th,
+                dlc_model_folder,
+                dlc_proc):
+    
     # 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:
-        keypts_xyp = dlc_live.get_pose(crop)  # third column is llk!
+        # scale crop here?
+        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
-        # add kpts of this crop to list
+        
+        #pdb.set_trace()
+        keypts_xyp[keypts_xyp[:,-1] < kpts_likelihood_th,:] = np_aux.fill(np.nan)
+        # add kpts of this crop to list 
         list_kpts_per_crop.append(keypts_xyp)
-
-    return list_kpts_per_crop
+        all_kypts.append(keypts_xyp)
+    
+    return list_kpts_per_crop

pre-requirements.txt

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

pyproject.toml

@@ -1,33 +0,0 @@
-[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

@@ -1,118 +0,0 @@
-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,13 +1,20 @@
-# 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
+ipython
+transformers
+Pillow
+gradio
+numpy
+torch
+torchvision
+timm
 seaborn
-deeplabcut[modelzoo,tf]==3.0.2
+deeplabcut[tf,modelzoo]>=3.0.0rc13
 deeplabcut-live
 ruamel.yaml==0.17.21
 dlclibrary
+argparse
 humanfriendly
+pandas
+statistics
+tqdm
 psutil
-ultralytics
+ultralytics

save_results.py

@@ -0,0 +1,56 @@
+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,228 +1,74 @@
 import gradio as gr
-from matplotlib import colormaps
-from matplotlib.colors import to_hex
-from PIL import Image
 
-from pytorch_utils import MAX_IMAGE_SIZE
-from viz_utils import COLORMAPS
-
-# shades built around the DeepLabCut docs palette (DeepLabCut/docs/_static/custom.css)
-DLC_PURPLE = gr.themes.Color(
-    name="dlc_purple",
-    c50="#f5edff",
-    c100="#ead7ff",
-    c200="#d9b8ff",
-    c300="#c084fc",
-    c400="#ac72f0",
-    c500="#9b5de5",
-    c600="#8550c4",
-    c700="#73439a",
-    c800="#5c357c",
-    c900="#4b236f",
-    c950="#2e1546",
-)
-DLC_TEAL = gr.themes.Color(
-    name="dlc_teal",
-    c50="#e8f8f7",
-    c100="#c9efec",
-    c200="#9fe2dd",
-    c300="#7ad3ce",
-    c400="#57c4be",
-    c500="#2fb0a8",
-    c600="#21a197",
-    c700="#1a8078",
-    c800="#16655f",
-    c900="#124f4a",
-    c950="#0a2e2b",
-)
-
-
-def dlc_theme():
-    # white text on purple 500 is 4.1:1, so filled buttons use 700 (7.0:1) and 600 (5.3:1)
-    return gr.themes.Default(primary_hue=DLC_PURPLE, secondary_hue=DLC_TEAL, neutral_hue="slate").set(
-        button_primary_background_fill="*primary_700",
-        button_primary_background_fill_hover="*primary_600",
-        button_primary_background_fill_dark="*primary_600",
-        button_primary_background_fill_hover_dark="*primary_500",
-        button_primary_text_color="white",
-        button_primary_text_color_dark="white",
-        button_primary_border_color="*primary_700",
-        button_primary_border_color_dark="*primary_600",
-    )
-
-
-def gradio_inputs_for_MD_DLC(backends_list, md_models_list, dlc_models_list):
+##############################
+def gradio_inputs_for_MD_DLC(md_models_list, # list(MD_models_dict.keys())
+                             dlc_models_list, # list(DLC_models_dict.keys())
+                             ):
     # 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 (TensorFlow legacy only)",
-        visible=gr_backend_input.value != "PyTorch",
-    )
+    gr_image_input = gr.inputs.Image(type="pil", label="Input Image")
 
-    gr_dlc_model_input = gr.Dropdown(
-        choices=dlc_models_list,
-        value="superanimal_quadruped",
-        type="value",
-        label="Select DeepLabCut model",
-    )
 
+    # Models
+    gr_mega_model_input = gr.inputs.Dropdown(choices=md_models_list,
+                                            default='md_v5a', # default option
+                                            type='value', # Type of value to be returned by component. "value" returns the string of the choice selected, "index" returns the index of the choice selected.
+                                            label='Select Detector model')
+    gr_dlc_model_input = gr.inputs.Dropdown(choices=dlc_models_list, # choices
+                                            default='superanimal_quadruped_dlcrnet', # default option
+                                            type='value', # Type of value to be returned by component. "value" returns the string of the choice selected, "index" returns the index of the choice selected.
+                                            label='Select DeepLabCut model')                                                                     
+    
     # Other inputs
-    gr_dlc_only_checkbox = gr.Checkbox(
-        value=False,
-        label="Run DeepLabCut only, directly on input image?",
-    )
-
-    # Gradio Slider signature is (minimum, maximum, value, step, ...)
-    gr_slider_conf_bboxes = gr.Slider(
-        minimum=0,
-        maximum=1,
-        value=0.2,
-        step=0.05,
-        label="Set confidence threshold for animal detections",
-    )
-
-    gr_slider_conf_keypoints = gr.Slider(
-        minimum=0,
-        maximum=1,
-        value=0.4,
-        step=0.05,
-        label="Set confidence threshold for keypoints",
-    )
-
-    # Data viz
-    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=False,
-            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,
-        gr_str_labels_checkbox,
-        gr_slider_conf_bboxes,
-        gr_slider_conf_keypoints,
-        gr_labels_font_style,
-        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>"""
-
-
+    gr_dlc_only_checkbox = gr.inputs.Checkbox(False,
+                                            label='Run DLClive only, directly on input image?')                                        
+    gr_str_labels_checkbox = gr.inputs.Checkbox(True,
+                                                label='Show bodypart labels?')   
+
+    gr_slider_conf_bboxes = gr.inputs.Slider(0,1,.05,0.2,
+                                            label='Set confidence threshold for animal detections')
+    gr_slider_conf_keypoints = gr.inputs.Slider(0,1,.05,0.4,
+                                                label='Set confidence threshold for keypoints')
+
+    # Data viz                                            
+    gr_keypt_color = gr.ColorPicker(value ="#862db7", label="choose color for keypoint label") 
+
+    gr_labels_font_style = gr.inputs.Dropdown(choices=['amiko', 'animals', 'nature', 'painter', 'zen'],
+                                            default='amiko',
+                                            type='value', 
+                                            label='Select keypoint label font')
+    gr_slider_font_size = gr.inputs.Slider(5,30,1,8,
+                                          label='Set font size')
+    gr_slider_marker_size = gr.inputs.Slider(1,20,1,9,
+                                            label='Set marker size')   
+
+    # list of inputs
+    return [gr_image_input,
+            gr_mega_model_input,
+            gr_dlc_model_input,
+            gr_dlc_only_checkbox,
+            gr_str_labels_checkbox,
+            gr_slider_conf_bboxes,
+            gr_slider_conf_keypoints,
+            gr_labels_font_style,
+            gr_slider_font_size, 
+            gr_keypt_color,
+            gr_slider_marker_size]     
+
+####################################################
 def gradio_outputs_for_MD_DLC():
-    gr_image_output = gr.Image(type="pil", label="Output Image")
-    gr_confidence_legend = gr.HTML("")
-    with gr.Row():
-        gr_file_download = gr.File(label="Download JSON file")
-        gr_image_download = gr.File(label="Download annotated image")
-    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)
-
-
+    # User interface: outputs
+    gr_image_output = gr.outputs.Image(type="pil", label="Output Image")
+    gr_file_download = gr.File(label="Download JSON file")
+    return [gr_image_output,
+            gr_file_download]
+
+##############################################
+# User interace: description
 def gradio_description_and_examples():
-    title = "DeepLabCut Model Zoo: SuperAnimals"
-    description = (
-        "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."
-    )
+    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>\."
+
 
-    examples = [
-        [image, "PyTorch", "md_v5a", "superanimal_quadruped", False, True, 0.5, 0.4, "amiko"]
-        + [font_size, "#ff0000", marker_size, False, "viridis", "#ff0000"]
-        for image in (
-            "examples/dog.jpeg",
-            "examples/cat.jpg",
-        )
-        for font_size, marker_size in [example_sizes(image)]
-    ]
+    examples = [['examples/dog.jpeg', 'md_v5a', 'superanimal_quadruped_dlcrnet', False, True, 0.5, 0.00, 'amiko',9, 'red', 3]]
 
-    return [title, description, examples]
+    return [title,description,examples]

viz_utils.py

@@ -1,38 +1,31 @@
-import json
-import math
-from datetime import date
-
+import json 
 import numpy as np
-from matplotlib import colormaps
-from PIL import ImageColor, ImageDraw, ImageFont
 
+from matplotlib import cm
+import matplotlib
+from PIL import Image, ImageColor, ImageFont, ImageDraw 
+import numpy as np
+import pdb
+from datetime import date
 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",
-}
-# perceptually uniform maps first; turbo separates neighbouring bodyparts best
-COLORMAPS = ["viridis", "plasma", "magma", "cividis", "turbo"]
-
+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,
-    colormap="viridis",
-):
+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,
+                            ):
     """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
@@ -46,232 +39,159 @@ def draw_keypoints_on_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]
-    confidences = [k[2] 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)
 
     # 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 = colormaps[colormap]
+    
+    #cmap = matplotlib.cm.get_cmap('hsv')
+    cmap2 = matplotlib.cm.get_cmap('Greys')
     # draw ellipses around keypoints
-    for i, (keypoint_x, keypoint_y) in enumerate(zip(keypoints_x, keypoints_y, strict=True)):
+    for i, (keypoint_x, keypoint_y) in enumerate(zip(keypoints_x, keypoints_y)):
+        round_fill = list(cm.viridis(norm(colores[i]),bytes=True))#[round(num*255) for num in list(cmap(i))[:3]] #check!
         # handling potential nans in the keypoints
         if np.isnan(keypoint_x).any():
             continue
-
-        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]
+        
+        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]
 
         # 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), #-------
-                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
-
+            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)
 
 #########################################
 #  Draw bboxes on image
-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 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
 
 ###########################################
-#  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"
-
+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'):
 
-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
+    """
+    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 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 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_only_dlc(dlc_outputs,map_label_id_to_str,model,output_file = 'dowload_predictions_dlc.json'):
 
-
-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)
     """
-    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())
-
-    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)
-
-
-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)
-
+    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))
 
-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).
+        kypts.append(aux1)
+    info['dlc_pred']  = dict(zip(labels,kypts))
 
-    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())
+    with open(output_file, 'w') as f:
+        json.dump(info, f, indent=1)
+        print('Output file saved at {}'.format(output_file))
 
-    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)
+    return output_file
 
 
-###########################################
+###########################################