Update app.py

app.py

@@ -20,6 +20,17 @@ from dlclibrary.dlcmodelzoo.modelzoo_download import (
     MODELOPTIONS,
 )
 
+
+
+# TESTING (passes) download the SuperAnimal models:
+#model = 'superanimal_topviewmouse'
+#train_dir = 'DLC_models/sa-tvm'
+#download_huggingface_model(model, train_dir)
+
+# grab demo data cooco cat:
+url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+image = Image.open(requests.get(url, stream=True).raw)
+
 # megadetector and dlc model look up
 MD_models_dict = {'md_v5a': "MD_models/md_v5a.0.0.pt", # 
                   'md_v5b': "MD_models/md_v5b.0.0.pt"}
@@ -29,15 +40,6 @@ DLC_models_dict = {'superanimal_topviewmouse': "DLC_models/sa-tvm",
                    'superanimal_quadreped': "DLC_models/sa-q",
                     'full_human': "DLC_models/DLC_human_dancing/"}
 
-# download the SuperAnimal models:
-model = 'superanimal_topviewmouse'
-train_dir = 'DLC_models/sa-tvm'
-download_huggingface_model(model, train_dir)
-
-# grab demo data cooco cat:
-url = "http://images.cocodataset.org/val2017/000000039769.jpg"
-image = Image.open(requests.get(url, stream=True).raw)
-
 
 #####################################################
 def predict_pipeline(img_input,
@@ -68,7 +70,6 @@ def predict_pipeline(img_input,
 
         ############################################################
 
-############################################################
     ## Get DLC model and label map  
     
     # If model is found: do not download (previous execution is likely within same day)
@@ -80,5 +81,10 @@ def predict_pipeline(img_input,
         path_to_DLCmodel = download_huggingface_model(dlc_model_input_str, 
                                          DLC_models_dict[dlc_model_input_str])
 
-    
-
+    # extract map label ids to strings
+    pose_cfg_path = os.path.join(DLC_models_dict[dlc_model_input_str],
+                                 'pose_cfg.yaml')
+    with open(pose_cfg_path, "r") as stream:
+        pose_cfg_dict = yaml.safe_load(stream) 
+    map_label_id_to_str = dict([(k,v) for k,v in zip([el[0] for el in pose_cfg_dict['all_joints']],  # pose_cfg_dict['all_joints'] is a list of one-element lists,
+                                                     pose_cfg_dict['all_joints_names'])])