adding sliders

app.py

@@ -60,7 +60,8 @@ def predict_md(im, size=640):
 
     return results #Image.fromarray(results.imgs[0]) ---return animals only?
 
-def crop_animal_detections(yolo_results, likelihood_th):
+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']
@@ -109,44 +110,54 @@ def crop_animal_detections(yolo_results, likelihood_th):
     # if detections_dict["category"] == "1":
     return list_np_animal_crops
 
-def predict_dlc(list_np_crops,DLCmodel,dlc_proc):
+def predict_dlc(list_np_crops,
+                kpts_likelihood_th,
+                DLCmodel, dlc_proc):
+
     # run dlc thru list of crops
     dlc_live = DLCLive(DLCmodel, processor=dlc_proc)
     dlc_live.init_inference(list_np_crops[0])
 
     list_kpts_per_crop = []
+    np_aux = np.empty((1,3)) # can I avoid hardcoding?
     for crop in list_np_crops:
-        # scale crop?
-        keypts = dlc_live.get_pose(crop) # third column is llk!
-        list_kpts_per_crop.append(keypts)
+        # 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_aux.fill(np.nan)
+        list_kpts_per_crop.append(keypts_xyp)
 
-        return list_kpts_per_crop
+    return list_kpts_per_crop
 
 
 
-def predict_pipeline(img_input):
+def predict_pipeline(img_input,
+                     bbox_likelihood_th,
+                     kpts_likelihood_th):
 
-    # these eventually user inputs....
+    # these will eventually be user inputs....
     path_to_DLCmodel =  "DLC_models/DLC_Cat_resnet_50_iteration-0_shuffle-0"
-    likelihood_th = 0.8
 
-    # Run Megadetector
+    ### Run Megadetector
     md_results = predict_md(img_input) #Image.fromarray(results.imgs[0])
 
     # Obtain animal crops with confidence above th
     list_crops = crop_animal_detections(md_results,
-                                        likelihood_th)
+                                        bbox_likelihood_th)
 
     # 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)):
 
@@ -158,18 +169,15 @@ def predict_pipeline(img_input):
                                 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_input.paste(img_crop,
-                        box = tuple([int(math.floor(t)) for t in md_results.xyxy[0][ic,:2]]))
-                                                                   
-        # plt.imshow(np_crop)
-        # plt.scatter(kpts_crop[:,0], kpts_crop[:,1], 40,
-        #             color='r')
-        # img_overlay = Image.frombytes('RGB', 
-        #                                fig.canvas.get_width_height(),
-        #                                fig.canvas.tostring_rgb())
-    return img_input #Image.fromarray(list_crops[0]) #Image.fromarray(md_results.imgs[0]) #
+        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 #
 
 
 ##########################################################
@@ -179,26 +187,28 @@ def predict_pipeline(img_input):
 MD_model = torch.hub.load('ultralytics/yolov5', 'custom', "model_weights/md_v5a.0.0.pt")
 
 
-
 ####################################################
 # Create user interface and launch
-#inputs = [image, chosen_model, size]
-inputs = gr.inputs.Image(type="pil", label="Input Image")
-outputs = gr.outputs.Image(type="pil", label="Output Image")
+gr_image_input = gr.inputs.Image(type="pil", label="Input Image")
+gr_image_output = gr.outputs.Image(type="pil", label="Output 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')
 #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
 
-title = "MegaDetector v5 + DLC live"
-description = "Detect and estimate pose of animals camera trap images using MegaDetector v5a + DLClive"
+gr_title = "MegaDetector v5 + DLClive"
+gr_description = "Detect and estimate the pose of animals in camera trap images, using MegaDetector v5a + DeepLabCut-live"
 # 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']]
 
 gr.Interface(predict_pipeline, 
-             inputs, 
-             outputs, 
-             title=title, 
-             description=description,
+             inputs=[gr_image_input,gr_slider_conf_bboxes,gr_slider_conf_keypoints], 
+             outputs=gr_image_output, 
+             title=gr_title, 
+             description=gr_description,
              theme="huggingface").launch(enable_queue=True)