added selectable dlc model, added option to run dlc model only on input image (mostly for testing)

app.py

@@ -112,7 +112,8 @@ def crop_animal_detections(yolo_results,
 
 def predict_dlc(list_np_crops,
                 kpts_likelihood_th,
-                DLCmodel, dlc_proc):
+                DLCmodel,
+                dlc_proc):
 
     # run dlc thru list of crops
     dlc_live = DLCLive(DLCmodel, processor=dlc_proc)
@@ -132,13 +133,17 @@ def predict_dlc(list_np_crops,
 
 
 def predict_pipeline(img_input,
+                     model_input_str,
+                     flag_dlc_only,
                      bbox_likelihood_th,
                      kpts_likelihood_th):
 
-    # these will eventually be user inputs....
-    path_to_DLCmodel =  "DLC_models/DLC_Cat_resnet_50_iteration-0_shuffle-0"
+    if model_input_str == 'full_cat':
+        path_to_DLCmodel =  "DLC_models/DLC_Cat_resnet_50_iteration-0_shuffle-0"
+    elif model_input_str == 'full_dog':
+        path_to_DLCmodel =  "DLC_models/DLC_Dog_resnet_50_iteration-0_shuffle-0"
 
-    ### Run Megadetector
+    # ### Run Megadetector
     md_results = predict_md(img_input) #Image.fromarray(results.imgs[0])
 
     # Obtain animal crops with confidence above th
@@ -148,36 +153,46 @@ def predict_pipeline(img_input,
     # Run DLC
     # TODO: add llk threshold for kpts too?
     dlc_proc = Processor()
-    list_kpts_per_crop = predict_dlc(list_crops,
-                                     kpts_likelihood_th,
-                                     path_to_DLCmodel,
-                                     dlc_proc)
-    
-
-    # # Produce final image
-    # fig = plt.Figure(md_results.imgs[0].shape[:2]) #figsize=(10,10)) #md_results.imgs[0].shape)
-    # list_annotated_crops = []
-    img_background = Image.fromarray(md_results.imgs[0])
-    for ic, (np_crop, kpts_crop) in enumerate(zip(list_crops,
-                                                  list_kpts_per_crop)):
-
-        ## Draw keypts on crop
-        img_crop = Image.fromarray(np_crop)
-        draw_keypoints_on_image(img_crop,
-                                kpts_crop, # a numpy array with shape [num_keypoints, 2].
+    if flag_dlc_only:
+        # compute kpts on input img
+        list_kpts_per_crop = predict_dlc([np.asarray(img_input)],#list_crops,--------
+                                         kpts_likelihood_th,
+                                         path_to_DLCmodel,
+                                         dlc_proc)
+        # draw kpts on input img
+        draw_keypoints_on_image(img_input,
+                                list_kpts_per_crop[0], # a numpy array with shape [num_keypoints, 2].
                                 color='red',
                                 radius=2,
-                                use_normalized_coordinates=False) # if True, then I should use md_results.xyxyn
-
-
-        # list_annotated_crops.append(img_crop)
-
-        ## Paste crop in original image
-        # https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.paste
-        img_background.paste(img_crop,
-                             box = tuple([int(math.floor(t)) for t in md_results.xyxy[0][ic,:2]]))
-                                                             
-    return img_background #Image.fromarray(list_crops[0]) #Image.fromarray(md_results.imgs[0]) #list_annotated_crops #
+                                use_normalized_coordinates=False)
+        return img_input
+
+    else:
+        # Compute kpts for each crop
+        list_kpts_per_crop = predict_dlc(list_crops,
+                                         kpts_likelihood_th,
+                                         path_to_DLCmodel,
+                                         dlc_proc)
+
+        # Produce final image
+        img_background = Image.fromarray(md_results.imgs[0]) # img_input?
+        for ic, (np_crop, kpts_crop) in enumerate(zip(list_crops,
+                                                      list_kpts_per_crop)):
+
+            ## Draw keypts on crop
+            img_crop = Image.fromarray(np_crop)
+            draw_keypoints_on_image(img_crop,
+                                    kpts_crop, # a numpy array with shape [num_keypoints, 2].
+                                    color='red',
+                                    radius=2,
+                                    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
+            img_background.paste(img_crop,
+                                box = tuple([int(math.floor(t)) for t in md_results.xyxy[0][ic,:2]]))
+                                                                
+        return img_background #Image.fromarray(list_crops[0]) #Image.fromarray(md_results.imgs[0]) #list_annotated_crops #
 
 
 ##########################################################
@@ -191,8 +206,15 @@ MD_model = torch.hub.load('ultralytics/yolov5', 'custom', "model_weights/md_v5a.
 # Create user interface and launch
 gr_image_input = gr.inputs.Image(type="pil", label="Input Image")
 gr_image_output = gr.outputs.Image(type="pil", label="Output Image")
+gr_dlc_model_input = gr.inputs.Dropdown(choices=['full_cat','full_dog'], # 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')
+                                        label='Set confidence threshold for animal detections')
 gr_slider_conf_keypoints = gr.inputs.Slider(0,1,.05,0,
                                              label='Set confidence threshold for keypoints')
 #image = gr.inputs.Image(type="pil", label="Input Image")
@@ -205,7 +227,11 @@ gr_description = "Detect and estimate the pose of animals in camera trap images,
 # examples = [['data/Macropod.jpg'], ['data/koala2.jpg'],['data/cat.jpg'],['data/BrushtailPossum.jpg']]
 
 gr.Interface(predict_pipeline, 
-             inputs=[gr_image_input,gr_slider_conf_bboxes,gr_slider_conf_keypoints], 
+             inputs=[gr_image_input,
+                     gr_dlc_model_input,
+                     gr_dlc_only_checkbox,
+                     gr_slider_conf_bboxes,
+                     gr_slider_conf_keypoints], 
              outputs=gr_image_output, 
              title=gr_title, 
              description=gr_description,