mwmathis commited on
Commit
a8badba
Β·
1 Parent(s): b11b05c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +18 -12
app.py CHANGED
@@ -20,6 +20,17 @@ from dlclibrary.dlcmodelzoo.modelzoo_download import (
20
  MODELOPTIONS,
21
  )
22
 
 
 
 
 
 
 
 
 
 
 
 
23
  # megadetector and dlc model look up
24
  MD_models_dict = {'md_v5a': "MD_models/md_v5a.0.0.pt", #
25
  'md_v5b': "MD_models/md_v5b.0.0.pt"}
@@ -29,15 +40,6 @@ DLC_models_dict = {'superanimal_topviewmouse': "DLC_models/sa-tvm",
29
  'superanimal_quadreped': "DLC_models/sa-q",
30
  'full_human': "DLC_models/DLC_human_dancing/"}
31
 
32
- # download the SuperAnimal models:
33
- model = 'superanimal_topviewmouse'
34
- train_dir = 'DLC_models/sa-tvm'
35
- download_huggingface_model(model, train_dir)
36
-
37
- # grab demo data cooco cat:
38
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
39
- image = Image.open(requests.get(url, stream=True).raw)
40
-
41
 
42
  #####################################################
43
  def predict_pipeline(img_input,
@@ -68,7 +70,6 @@ def predict_pipeline(img_input,
68
 
69
  ############################################################
70
 
71
- ############################################################
72
  ## Get DLC model and label map
73
 
74
  # If model is found: do not download (previous execution is likely within same day)
@@ -80,5 +81,10 @@ def predict_pipeline(img_input,
80
  path_to_DLCmodel = download_huggingface_model(dlc_model_input_str,
81
  DLC_models_dict[dlc_model_input_str])
82
 
83
-
84
-
 
 
 
 
 
 
20
  MODELOPTIONS,
21
  )
22
 
23
+
24
+
25
+ # TESTING (passes) download the SuperAnimal models:
26
+ #model = 'superanimal_topviewmouse'
27
+ #train_dir = 'DLC_models/sa-tvm'
28
+ #download_huggingface_model(model, train_dir)
29
+
30
+ # grab demo data cooco cat:
31
+ url = "http://images.cocodataset.org/val2017/000000039769.jpg"
32
+ image = Image.open(requests.get(url, stream=True).raw)
33
+
34
  # megadetector and dlc model look up
35
  MD_models_dict = {'md_v5a': "MD_models/md_v5a.0.0.pt", #
36
  'md_v5b': "MD_models/md_v5b.0.0.pt"}
 
40
  'superanimal_quadreped': "DLC_models/sa-q",
41
  'full_human': "DLC_models/DLC_human_dancing/"}
42
 
 
 
 
 
 
 
 
 
 
43
 
44
  #####################################################
45
  def predict_pipeline(img_input,
 
70
 
71
  ############################################################
72
 
 
73
  ## Get DLC model and label map
74
 
75
  # If model is found: do not download (previous execution is likely within same day)
 
81
  path_to_DLCmodel = download_huggingface_model(dlc_model_input_str,
82
  DLC_models_dict[dlc_model_input_str])
83
 
84
+ # extract map label ids to strings
85
+ pose_cfg_path = os.path.join(DLC_models_dict[dlc_model_input_str],
86
+ 'pose_cfg.yaml')
87
+ with open(pose_cfg_path, "r") as stream:
88
+ pose_cfg_dict = yaml.safe_load(stream)
89
+ 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,
90
+ pose_cfg_dict['all_joints_names'])])