Update app.py
Browse files
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'])])
|