added font, font color, pose size,

app.py

@@ -31,13 +31,61 @@ FONTS = {'amiko': "font/Amiko-Regular.ttf",
         'zen': "font/ZEN.TTF"}
 
 
+#############################################
+# %% 
+# Create user interface and launch: all inputs
+gr_image_input = gr.inputs.Image(type="pil", label="Input Image")
+
+
+gr_dlc_model_input = gr.inputs.Dropdown(choices=['full_cat','full_dog', 'monkey_face', 'full_human', 'full_monkey'], # choices
+                                        default='full_cat', # 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 DLC model')
+gr_dlc_only_checkbox = gr.inputs.Checkbox(False,
+                                         label='Run DLClive only, directly on input image?')                                        
+
+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')
+                                            
+gr_keypt_color = gr.ColorPicker(label="choose colour for pose") 
+
+gr_pose_font_input = gr.inputs.Dropdown(choices=['amiko', 'nature', 'painter', 'animals', 'zen'],
+                                        default='amiko',
+                                        type='value', 
+                                        label='Select pose font')
+gr_slider_font_size = gr.inputs.Slider(5,30,1,8,
+                                        label='Set pose font size')
+
+gr_slider_pose_size = gr.inputs.Slider(1,20,1,2,
+                                        label='Set pose size')                                          
+
+inputs = [gr_image_input,
+                     gr_dlc_model_input,
+                     gr_dlc_only_checkbox,
+                     gr_slider_conf_bboxes,
+                     gr_slider_conf_keypoints,
+                     gr_pose_font_input,
+                     gr_slider_font_size, 
+                     gr_keypt_color,
+                     gr_slider_pose_size,
+                     ]                                                                                
+
+#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
+
+#########################################
+# %% 
 def draw_keypoints_on_image(image,
                             keypoints,
                             map_label_id_to_str,
-                            color='red',
-                            radius=2,
                             use_normalized_coordinates=True,
                             gr_pose_font_input='amiko',
+                            gr_slider_font_size=8,
+                            gr_keypt_color="#ff0000",
+                            gr_slider_pose_size='2'
                             ):
     """Draws keypoints on an image.
     Modified from:
@@ -64,21 +112,21 @@ def draw_keypoints_on_image(image,
         keypoints_y = tuple([im_height * y for y in keypoints_y])
 
     # draw ellipses around keypoints and add string labels
-    font = ImageFont.truetype(FONTS[gr_pose_font_input], 8) # font = ImageFont.truetype(<font-file>, <font-size>)
+    font = ImageFont.truetype(FONTS[gr_pose_font_input], gr_slider_font_size) # 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)
+        draw.ellipse([(keypoint_x - gr_slider_pose_size, keypoint_y - gr_slider_pose_size),
+                      (keypoint_x + gr_slider_pose_size, keypoint_y + gr_slider_pose_size)],
+                       outline=gr_keypt_color, fill=gr_keypt_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), #-------
+        draw.text((keypoint_x + gr_slider_pose_size, keypoint_y + gr_slider_pose_size),#(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
 def predict_md(im, size=640):
     # resize image
@@ -91,6 +139,7 @@ def predict_md(im, size=640):
 
     return results #Image.fromarray(results.imgs[0]) ---return animals only?
 
+# %% 
 def crop_animal_detections(yolo_results, 
                            likelihood_th):
     ## crop if animal and return list of crops
@@ -140,6 +189,7 @@ def crop_animal_detections(yolo_results,
     # if detections_dict["category"] == "1":
     return list_np_animal_crops
 
+# %%
 def predict_dlc(list_np_crops,
                 kpts_likelihood_th,
                 DLCmodel,
@@ -162,13 +212,16 @@ def predict_dlc(list_np_crops,
     return list_kpts_per_crop
 
 
-
+# %% 
 def predict_pipeline(img_input,
                      model_input_str,
                      flag_dlc_only,
                      bbox_likelihood_th,
                      kpts_likelihood_th,
-                     gr_pose_font_input,
+                     gr_pose_font_input=gr_pose_font_input,
+                     gr_slider_font_size=gr_slider_font_size,
+                     gr_keypt_color=gr_keypt_color,
+                     gr_slider_pose_size=gr_slider_pose_size,
                      ):
 
     ############################################################
@@ -188,6 +241,7 @@ def predict_pipeline(img_input,
     elif model_input_str == 'full_monkey':
         path_to_DLCmodel =  "model/DLC_monkey_resnet_50_iteration-0_shuffle-1"
         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:
@@ -219,10 +273,11 @@ def predict_pipeline(img_input,
         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,
-                                gr_pose_font_input=gr_pose_font_input
+                                gr_pose_font_input=gr_pose_font_input,
+                                gr_slider_font_size=gr_slider_font_size,
+                                gr_keypt_color=gr_keypt_color,
+                                gr_slider_pose_size=gr_slider_pose_size,
                                 )
         return img_input, #list_kpts_per_crop
 
@@ -243,10 +298,11 @@ def predict_pipeline(img_input,
             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
                                     gr_pose_font_input=gr_pose_font_input,
+                                    gr_slider_font_size=gr_slider_font_size,
+                                    gr_keypt_color=gr_keypt_color,
+                                    gr_slider_pose_size=gr_slider_pose_size,
                                     )
 
             ## Paste crop in original image
@@ -258,53 +314,16 @@ def predict_pipeline(img_input,
 
 
 ##########################################################
+# %% 
 # Get MegaDetector model
 # TODO: Allow user selectable model?
 # models = ["model_weights/md_v5a.0.0.pt","model_weights/md_v5b.0.0.pt"]
 MD_model = torch.hub.load('ultralytics/yolov5', 'custom', "megadet_model/md_v5b.0.0.pt")
 
 
-####################################################
-# Create user interface and launch: all inputs
-gr_image_input = gr.inputs.Image(type="pil", label="Input Image")
-
-
-gr_dlc_model_input = gr.inputs.Dropdown(choices=['full_cat','full_dog', 'monkey_face', 'full_human', 'full_monkey'], # choices
-                                        default='full_cat', # 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 DLC model')
-gr_dlc_only_checkbox = gr.inputs.Checkbox(False,
-                                         label='Run DLClive only, directly on input image?')                                        
-
-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')
-                                            
-gr_keypt_color = gr.ColorPicker(label="choose colour for pose") 
-
-gr_pose_font_input = gr.inputs.Dropdown(choices=['amiko', 'nature', 'painter', 'animals', 'zen'],
-                                        default='amiko',
-                                        type='value', 
-                                        label='Select pose font')
-gr_slider_font_size = gr.inputs.Slider(5,30,1,8,
-                                        label='Set pose font size')                                        
-
-inputs = [gr_image_input,
-                     gr_dlc_model_input,
-                     gr_dlc_only_checkbox,
-                     gr_slider_conf_bboxes,
-                     gr_slider_conf_keypoints,
-                     #gr_keypt_color,
-                     gr_pose_font_input,
-                     #gr_slider_font_size, 
-                     ]                                                                                
-
-#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
 
 ####################################################
+# %% 
 # Create user interface and launch: all outputs
 gr_gallery_output = gr.Gallery(type="pil", label="Output Gallery")
 gr_json_output = gr.JSON(label='megadetector json')
@@ -316,12 +335,15 @@ outputs = [gr_gallery_output, #gr_json_output,
 #gr_pose_output,
 ]
 
-gr_title = "MegaDetector v5 + DLClive"
+gr_title = "MegaDetector v5 + DLClive by Sofi, Nesli, Nirel, Vic"
 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>"
+                  Builds up on work from <a href='https://huggingface.co/spaces/hlydecker/MegaDetector_v5'>hlydecker/MegaDetector_v5</a> \
+                  <a href='https://huggingface.co/spaces/sofmi/MegaDetector_DLClive</a> \
+                  <a href='https://huggingface.co/spaces/Neslihan/megadetector_dlcmodels</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']]
 
+# %% 
 demo = gr.Interface(predict_pipeline, 
              inputs=inputs,
              outputs=outputs, 
@@ -331,22 +353,5 @@ demo = gr.Interface(predict_pipeline,
              #live=True
              )
 
-demo.launch(enable_queue=True)
-
-# def dlclive_pose(model, crop_np, crop, fname, index,dlc_proc):
-#     dlc_live = DLCLive(model, processor=dlc_proc) 
-#     dlc_live.init_inference(crop_np)
-#     keypts = dlc_live.get_pose(crop_np) 
-#     savetxt(str(fname)+ '_' + str(index) + '.csv' , keypts, delimiter=',')
-#     xpose = []
-#     ypose = []
-#     for key in keypts[:,2]:
-#        # if key > 0.05: # which value do we need here?
-#             i = np.where(keypts[:,2]==key)
-#             xpose.append(keypts[i,0])
-#             ypose.append(keypts[i,1])
-#     plt.imshow(crop)
-#     plt.scatter(xpose[:], ypose[:], 40, color='cyan')
-#     plt.savefig(str(fname)+ '_' + str(index) + '.png')
-#     plt.show()
-#     plt.clf()
+demo.launch(enable_queue=True, share=True)
+