added megadet model & bodypart colour

app.py

@@ -30,7 +30,9 @@ FONTS = {'amiko': "font/Amiko-Regular.ttf",
         'animals': "font/UncialAnimals.ttf", 
         'zen': "font/ZEN.TTF"}
 
-
+Megadet_Models = {'md_v5a': "megadet_model/md_v5a.0.0.pt",
+                  'md_v5b': "megadet_model/md_v5b.0.0.pt"}
+                  
 #############################################
 # %% 
 # Create user interface and launch: all inputs
@@ -60,7 +62,10 @@ gr_slider_font_size = gr.inputs.Slider(5,30,1,8,
 
 gr_slider_pose_size = gr.inputs.Slider(1,20,1,2,
                                         label='Set pose size')                                          
-
+gr_mega_model_input = gr.inputs.Dropdown(choices=['md_v5a','md_v5a'],
+                                         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 Megadetector model')
 inputs = [gr_image_input,
                      gr_dlc_model_input,
                      gr_dlc_only_checkbox,
@@ -70,6 +75,7 @@ inputs = [gr_image_input,
                      gr_slider_font_size, 
                      gr_keypt_color,
                      gr_slider_pose_size,
+                     gr_mega_model_input,
                      ]                                                                                
 
 #image = gr.inputs.Image(type="pil", label="Input Image")
@@ -122,17 +128,17 @@ def draw_keypoints_on_image(image,
         # draw.text((x, y),"Sample Text",(r,g,b))
         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
+                  (gr_keypt_color), # rgb
                   font=font)
                     
 ############################################
 # %% 
 # Predict detections with MegaDetector v5a model
-def predict_md(im, size=640):
+def predict_md(im, size=640, gr_mega_model_input='md_v5a'):
     # resize image
     g = (size / max(im.size))  # gain
     im = im.resize((int(x * g) for x in im.size), Image.ANTIALIAS)  # resize
-    
+    MD_model = torch.hub.load('ultralytics/yolov5', 'custom', Megadet_Models[gr_mega_model_input])
     ## detect objects
     results = MD_model(im)  # inference # vars(results).keys()= dict_keys(['imgs', 'pred', 'names', 'files', 'times', 'xyxy', 'xywh', 'xyxyn', 'xywhn', 'n', 't', 's'])
     results.render()  # updates results.imgs with boxes and labels
@@ -222,6 +228,7 @@ def predict_pipeline(img_input,
                      gr_slider_font_size=gr_slider_font_size,
                      gr_keypt_color=gr_keypt_color,
                      gr_slider_pose_size=gr_slider_pose_size,
+                     gr_mega_model_input=gr_mega_model_input,
                      ):
 
     ############################################################
@@ -251,7 +258,7 @@ def predict_pipeline(img_input,
 
     ############################################################                                               
     # ### Run Megadetector
-    md_results = predict_md(img_input) #Image.fromarray(results.imgs[0])
+    md_results = predict_md(img_input, gr_mega_model_input=gr_mega_model_input) #Image.fromarray(results.imgs[0])
 
     ################################################################
     # Obtain animal crops for bboxes with confidence above th
@@ -318,7 +325,7 @@ 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")
+#MD_model = torch.hub.load('ultralytics/yolov5', 'custom', "megadet_model/md_v5b.0.0.pt")