added string labels, reviewing kpts

app.py

@@ -1,51 +1,68 @@
-
 import gradio as gr
+
 import torch
 import torchvision
-import numpy as np
+from dlclive import DLCLive, Processor
+
 from PIL import Image
-import PIL.ImageDraw as ImageDraw
-import math
-import pdb
+from PIL import ImageFont
+from PIL import ImageDraw 
 
-from dlclive import DLCLive, Processor
+import numpy as np
+import math
+# import json
+import os
+import yaml
 
-import matplotlib.pyplot as plt
+import pdb
 
 #########################################
-# https://www.programcreek.com/python/?code=fjchange%2Fobject_centric_VAD%2Fobject_centric_VAD-master%2Fobject_detection%2Futils%2Fvisualization_utils.py
+
 def draw_keypoints_on_image(image,
                             keypoints,
+                            map_label_id_to_str,
                             color='red',
                             radius=2,
-                            use_normalized_coordinates=True):
-  """Draws keypoints on an image.
-
-  Args:
+                            use_normalized_coordinates=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
+    Args:
     image: a PIL.Image object.
     keypoints: a numpy array with shape [num_keypoints, 2].
     color: color to draw the keypoints with. Default is red.
     radius: keypoint radius. Default value is 2.
     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)
-
-  im_width, im_height = image.size
-  keypoints_x = [k[1] for k in keypoints]
-  keypoints_y = [k[0] for k in keypoints]
-
-  # adjust keypoints coords if required
-  if use_normalized_coordinates:
-    keypoints_x = tuple([im_width * x for x in keypoints_x])
-    keypoints_y = tuple([im_height * y for y in keypoints_y])
-
-  # draw ellipses around keypoints
-  for keypoint_x, keypoint_y in zip(keypoints_x, keypoints_y):
-    draw.ellipse([(keypoint_x - radius, keypoint_y - radius),
-                  (keypoint_x + radius, keypoint_y + radius)],
-                  outline=color, fill=color)
+        relative to the image.  Otherwise treat them as absolute.
+    """
+    # get a drawing context
+    draw = ImageDraw.Draw(image)
+    # font = ImageFont.truetype("sans-serif.ttf", 16)
+
+    im_width, im_height = image.size
+    keypoints_x = [k[1] for k in keypoints]
+    keypoints_y = [k[0] for k in keypoints]
+
+    # adjust keypoints coords if required
+    if use_normalized_coordinates:
+        keypoints_x = tuple([im_width * x for x in keypoints_x])
+        keypoints_y = tuple([im_height * y for y in keypoints_y])
+
+    # draw ellipses around keypoints and add string labels
+    font = ImageFont.truetype("Amiko-Regular.ttf", 8) # font = ImageFont.truetype(<font-file>, <font-size>)
+    for i, (keypoint_x, keypoint_y) in enumerate(zip(keypoints_x, keypoints_y)):
+        draw.ellipse([(keypoint_x - radius, keypoint_y - radius),
+                      (keypoint_x + radius, keypoint_y + radius)],
+                       outline=color, fill=color)
+
+        # add string labels around keypoints
+        # draw.text((x, y),"Sample Text",(r,g,b))
+        draw.text((keypoint_x + radius, keypoint_y + radius),#(0.5*im_width, 0.5*im_height), #-------
+                  map_label_id_to_str[i],#"Sample Text",
+                  (255,0,0), # rgb
+                  font=font)
+                    
 ############################################
 
 # Predict detections with MegaDetector v5a model
@@ -63,7 +80,6 @@ def predict_md(im, size=640):
 def crop_animal_detections(yolo_results, 
                            likelihood_th):
     ## crop if animal and return list of crops
-
     list_labels_as_str = yolo_results.names #['animal', 'person', 'vehicle']
     list_np_animal_crops = []
 
@@ -124,8 +140,15 @@ def predict_dlc(list_np_crops,
     for crop in list_np_crops:
         # scale crop here?
         keypts_xyp = dlc_live.get_pose(crop) # third column is llk!
+        # ATT! coord syst for keypoints is bottom left corner? change here to top left for PIL coord system
+        pdb.set_trace()
+        keypts_xyp = np.column_stack((crop.shape[1]-keypts_xyp[:,0],
+                                      crop.shape[0]-keypts_xyp[:,1],
+                                      keypts_xyp[:,2]))
+        pdb.set_trace()
         # set kpts below threhsold to nan
         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
@@ -138,21 +161,35 @@ def predict_pipeline(img_input,
                      bbox_likelihood_th,
                      kpts_likelihood_th):
 
+    ############################################################
+    ## Get DLC model and labels as strings   
     if model_input_str == 'full_cat':
         path_to_DLCmodel =  "DLC_models/DLC_Cat_resnet_50_iteration-0_shuffle-0"
+        pose_cfg_path = os.path.join(path_to_DLCmodel,'pose_cfg.yaml')
     elif model_input_str == 'full_dog':
         path_to_DLCmodel =  "DLC_models/DLC_Dog_resnet_50_iteration-0_shuffle-0"
+        pose_cfg_path = os.path.join(path_to_DLCmodel,'pose_cfg.yaml')
 
+    # read pose cfg as dict
+    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 Megadetector
     md_results = predict_md(img_input) #Image.fromarray(results.imgs[0])
 
-    # Obtain animal crops with confidence above th
+    ################################################################
+    # Obtain animal crops for bboxes with confidence above th
     list_crops = crop_animal_detections(md_results,
                                         bbox_likelihood_th)
 
+    ##############################################################
     # Run DLC
-    # TODO: add llk threshold for kpts too?
     dlc_proc = Processor()
+
+    # 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)],#list_crops,--------
@@ -162,6 +199,7 @@ def predict_pipeline(img_input,
         # draw kpts on input img
         draw_keypoints_on_image(img_input,
                                 list_kpts_per_crop[0], # a numpy array with shape [num_keypoints, 2].
+                                map_label_id_to_str,
                                 color='red',
                                 radius=2,
                                 use_normalized_coordinates=False)
@@ -183,9 +221,11 @@ def predict_pipeline(img_input,
             img_crop = Image.fromarray(np_crop)
             draw_keypoints_on_image(img_crop,
                                     kpts_crop, # a numpy array with shape [num_keypoints, 2].
+                                    map_label_id_to_str,
                                     color='red',
                                     radius=2,
-                                    use_normalized_coordinates=False) # if True, then I should use md_results.xyxyn
+                                    use_normalized_coordinates=False,  # if True, then I should use md_results.xyxyn
+                                    )
 
             ## Paste crop in original image
             # https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.paste
@@ -216,13 +256,14 @@ gr_dlc_only_checkbox = gr.inputs.Checkbox(False,
 gr_slider_conf_bboxes = gr.inputs.Slider(0,1,.05,0.8,
                                         label='Set confidence threshold for animal detections')
 gr_slider_conf_keypoints = gr.inputs.Slider(0,1,.05,0,
-                                             label='Set confidence threshold for keypoints')
+                                            label='Set confidence threshold for keypoints')
 #image = gr.inputs.Image(type="pil", label="Input Image")
 #chosen_model = gr.inputs.Dropdown(choices = models, value = "model_weights/md_v5a.0.0.pt",type = "value", label="Model Weight")
 #size = 640
 
 gr_title = "MegaDetector v5 + DLClive"
-gr_description = "Detect and estimate the pose of animals in camera trap images, using MegaDetector v5a + DeepLabCut-live"
+gr_description = "Detect and estimate the pose of animals in camera trap images, using MegaDetector v5a + DeepLabCut-live. \
+                  Builds up on work from <a href='https://huggingface.co/spaces/hlydecker/MegaDetector_v5'>hlydecker/MegaDetector_v5</a>"
 # article = "<p style='text-align: center'>This app makes predictions using a YOLOv5x6 model that was trained to detect animals, humans, and vehicles in camera trap images; find out more about the project on <a href='https://github.com/microsoft/CameraTraps'>GitHub</a>. This app was built by Henry Lydecker but really depends on code and models developed by <a href='http://ecologize.org/'>Ecologize</a> and <a href='http://aka.ms/aiforearth'>Microsoft AI for Earth</a>. Find out more about the YOLO model from the original creator, <a href='https://pjreddie.com/darknet/yolo/'>Joseph Redmon</a>. YOLOv5 is a family of compound-scaled object detection models trained on the COCO dataset and developed by Ultralytics, and includes simple functionality for Test Time Augmentation (TTA), model ensembling, hyperparameter evolution, and export to ONNX, CoreML and TFLite. <a href='https://github.com/ultralytics/yolov5'>Source code</a> | <a href='https://pytorch.org/hub/ultralytics_yolov5'>PyTorch Hub</a></p>"
 # examples = [['data/Macropod.jpg'], ['data/koala2.jpg'],['data/cat.jpg'],['data/BrushtailPossum.jpg']]
 

requirements.txt

@@ -4,4 +4,5 @@ torchvision
 numpy
 opencv-python
 seaborn
-deeplabcut-live
+deeplabcut-live
+pyyaml