Update app.py

app.py

@@ -90,6 +90,78 @@ def predict_pipeline(img_input,
                                                      pose_cfg_dict['all_joints_names'])])
 
 
+##############################################################
+    # Run DLC and visualize results
+    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)],
+                                         kpts_likelihood_th,
+                                         path_to_DLCmodel,
+                                         dlc_proc)
+        # draw kpts on input img #fix!
+        draw_keypoints_on_image(img_input,
+                                list_kpts_per_crop[0], # a numpy array with shape [num_keypoints, 2].
+                                map_label_id_to_str,
+                                flag_show_str_labels,
+                                use_normalized_coordinates=False,
+                                font_style=font_style,
+                                font_size=font_size,
+                                keypt_color=keypt_color,
+                                marker_size=marker_size)
+
+        donw_file = save_results_only_dlc(list_kpts_per_crop[0], map_label_id_to_str,dlc_model_input_str)
+
+        return img_input, donw_file
+
+    else:
+        # Compute kpts for each crop
+        list_kpts_per_crop = predict_dlc(list_crops,
+                                         kpts_likelihood_th,
+                                         path_to_DLCmodel,
+                                         dlc_proc)
+        
+        # resize input image to match megadetector output
+        img_background = img_input.resize((md_results.ims[0].shape[1],
+                                           md_results.ims[0].shape[0]))
+        
+        # draw keypoints on each crop and paste to background img
+        for ic, (np_crop, kpts_crop) in enumerate(zip(list_crops,
+                                                      list_kpts_per_crop)):
+
+            img_crop = Image.fromarray(np_crop)
+
+            # Draw keypts on crop
+            draw_keypoints_on_image(img_crop,
+                                    kpts_crop, # a numpy array with shape [num_keypoints, 2].
+                                    map_label_id_to_str,
+                                    flag_show_str_labels,
+                                    use_normalized_coordinates=False,  # if True, then I should use md_results.xyxyn for list_kpts_crop
+                                    font_style=font_style,
+                                    font_size=font_size,
+                                    keypt_color=keypt_color,
+                                    marker_size=marker_size)
+
+            # Paste crop in original image
+            img_background.paste(img_crop, 
+                                 box = tuple([int(t) for t in md_results.xyxy[0][ic,:2]]))
+
+            # Plot bbox
+            bb_per_animal = md_results.xyxy[0].tolist()[ic]
+            pred = md_results.xyxy[0].tolist()[ic][4]
+            if bbox_likelihood_th < pred:
+                draw_bbox_w_text(img_background, 
+                                    bb_per_animal,
+                                    font_style=font_style,
+                                    font_size=font_size)  # TODO: add selectable color for bbox?
+
+            
+        # Save detection results as json
+        download_file  = save_results_as_json(md_results,list_kpts_per_crop,map_label_id_to_str, bbox_likelihood_th,dlc_model_input_str,mega_model_input)         
+
+        return img_background, download_file