added dict for dlc model (fix not falling into any if case)

app.py

@@ -30,7 +30,13 @@ FONTS = {'amiko': "font/Amiko-Regular.ttf",
 
 Megadet_Models = {'md_v5a': "megadet_model/md_v5a.0.0.pt",
                   'md_v5b': "megadet_model/md_v5b.0.0.pt"}
-                  
+
+DLC_models = {'full_cat': "model/DLC_Cat_resnet_50_iteration-0_shuffle-0",
+              'full_dog': "model/DLC_Dog_resnet_50_iteration-0_shuffle-0",
+              'primate_face': "model/DLC_FacialLandmarks_resnet_50_iteration-1_shuffle-1",
+              'full_human': "model/DLC_human_dancing_resnet_101_iteration-0_shuffle-1",
+              'full_macaque': 'model/DLC_monkey_resnet_50_iteration-0_shuffle-1'}                  
+           
 #########################################
 #  Draw keypoints on image
 def draw_keypoints_on_image(image,
@@ -164,7 +170,7 @@ def predict_dlc(list_np_crops,
 #####################################################
 def predict_pipeline(img_input,
                      mega_model_input,
-                     model_input_str,
+                     dlc_model_input_str,
                      flag_dlc_only,
                      flag_show_str_labels,
                      bbox_likelihood_th,
@@ -178,23 +184,9 @@ def predict_pipeline(img_input,
     ############################################################
     ## Get DLC model and labels as strings   
     # TODO: make a dict as for megadetector
-    if model_input_str == 'full_cat':
-        path_to_DLCmodel =  "model/DLC_Cat_resnet_50_iteration-0_shuffle-0"
-        pose_cfg_path = os.path.join(path_to_DLCmodel,'pose_cfg.yaml')
-    elif model_input_str == 'full_dog':
-        path_to_DLCmodel =  "model/DLC_Dog_resnet_50_iteration-0_shuffle-0"
-        pose_cfg_path = os.path.join(path_to_DLCmodel,'pose_cfg.yaml')
-    elif model_input_str == 'primate_face':
-        path_to_DLCmodel =  "model/DLC_FacialLandmarks_resnet_50_iteration-1_shuffle-1"
-        pose_cfg_path = os.path.join(path_to_DLCmodel,'pose_cfg.yaml')
-    elif model_input_str == 'full_human':
-        path_to_DLCmodel =  "model/DLC_human_dancing_resnet_101_iteration-0_shuffle-1"
-        pose_cfg_path = os.path.join(path_to_DLCmodel,'pose_cfg.yaml')
-    elif model_input_str == 'full_macaque':
-        path_to_DLCmodel =  "model/DLC_monkey_resnet_50_iteration-0_shuffle-1"
-        pose_cfg_path = os.path.join(path_to_DLCmodel,'pose_cfg.yaml')        
+    path_to_DLCmodel = DLC_models[dlc_model_input_str]
+    pose_cfg_path = os.path.join(path_to_DLCmodel,'pose_cfg.yaml')        
                      
-        
     # extract map label ids to strings
     # pose_cfg_dict['all_joints'] is a list of one-element lists,
     with open(pose_cfg_path, "r") as stream:
@@ -283,11 +275,11 @@ gr_image_input = gr.inputs.Image(type="pil", label="Input Image")
 
 
 # Models
-gr_dlc_model_input = gr.inputs.Dropdown(choices=['full_cat','full_dog', 'primate_face', 'full_human', 'full_macaque'], # choices
+gr_dlc_model_input = gr.inputs.Dropdown(choices=DLC_models.keys(), # 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 DeepLabCut model')                                                                     
-gr_mega_model_input = gr.inputs.Dropdown(choices=['md_v5a','md_v5b'],
+gr_mega_model_input = gr.inputs.Dropdown(choices=Megadet_Models.keys(),
                                          default='md_v5a', # 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 MegaDetector model')