save json to dlc only

app.py

@@ -13,7 +13,7 @@ from PIL import Image, ImageColor, ImageFont, ImageDraw
 from DLC_models.download_utils import DownloadModel
 from dlclive import DLCLive, Processor
 
-from viz_utils import save_results_as_json, draw_keypoints_on_image, draw_bbox_w_text
+from viz_utils import save_results_as_json, draw_keypoints_on_image, draw_bbox_w_text, save_results_only_dlc
 from detection_utils import predict_md, crop_animal_detections, predict_dlc
 from ui_utils import gradio_inputs_for_MD_DLC, gradio_outputs_for_MD_DLC, gradio_description_and_examples
 
@@ -105,8 +105,10 @@ def predict_pipeline(img_input,
                                 font_size=font_size,
                                 keypt_color=keypt_color,
                                 marker_size=marker_size)
-                                    
-        return img_input, []
+
+        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
@@ -151,10 +153,7 @@ def predict_pipeline(img_input,
 
             
         # 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)         
+        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
 

viz_utils.py

@@ -5,7 +5,9 @@ from matplotlib import cm
 import matplotlib
 from PIL import Image, ImageColor, ImageFont, ImageDraw 
 import numpy as np
-
+import pdb
+from datetime import date
+today = date.today()
 FONTS = {'amiko': "fonts/Amiko-Regular.ttf",
         'nature': "fonts/LoveNature.otf", 
         'painter':"fonts/PainterDecorator.otf",
@@ -104,11 +106,7 @@ def draw_bbox_w_text(img,
     return imgR
 
 ###########################################
-def save_results_as_json(md_results, 
-                            dlc_outputs,
-                            map_dlc_label_id_to_str,
-                            thr,
-                            path_to_output_file = 'download_predictions.json'):
+def save_results_as_json(md_results, dlc_outputs, map_dlc_label_id_to_str, thr,model,mega_model_input, path_to_output_file = 'download_predictions.json'):
 
     """
     Output detections as json file
@@ -116,6 +114,8 @@ def save_results_as_json(md_results,
     """
     # initialise dict to save to json
     info = {}
+    info['date'] = str(today)
+    info['MD_model'] = str(mega_model_input)
     # info from megaDetector
     info['file']= md_results.files[0]
     number_bb = len(md_results.xyxy[0].tolist())
@@ -143,6 +143,7 @@ def save_results_as_json(md_results,
         aux['confidence MD'] = confidence
         
         # DLC output
+        info['dlc_model'] = model
         kypts = []
         for s in dlc_outputs[i]:
             aux1 = []
@@ -161,5 +162,29 @@ def save_results_as_json(md_results,
     return path_to_output_file
 
 
+def save_results_only_dlc(dlc_outputs,map_label_id_to_str,model,output_file = 'dowload_predictions_dlc.json'):
+
+    """
+    write json dlc output
+    """
+    info = {}
+    info['date'] = today
+    labels = [n for n in map_label_id_to_str.values()]
+    info['dlc_model'] = model
+    kypts = []
+    for s in dlc_outputs:
+        aux1 = []
+        for j in s:
+            aux1.append(float(j))
+
+        kypts.append(aux1)
+    info['dlc_pred']  = dict(zip(labels,kypts))
+
+    with open(output_file, 'w') as f:
+        json.dump(info, f, indent=1)
+        print('Output file saved at {}'.format(output_file))
+
+    return output_file
+
 
 ###########################################