Spaces:
Build error
Build error
sabrinabenas
commited on
Commit
β’
39504e9
1
Parent(s):
f52e6b1
add save results
Browse files- app.py +29 -32
- save_results.py +59 -0
app.py
CHANGED
@@ -19,7 +19,7 @@ import math
|
|
19 |
import os
|
20 |
import yaml
|
21 |
from model.models import DownloadModel
|
22 |
-
|
23 |
import pdb
|
24 |
|
25 |
#########################################
|
@@ -112,8 +112,7 @@ def predict_md(im,
|
|
112 |
|
113 |
## detect objects
|
114 |
results = MD_model(im) # inference # vars(results).keys()= dict_keys(['imgs', 'pred', 'names', 'files', 'times', 'xyxy', 'xywh', 'xyxyn', 'xywhn', 'n', 't', 's'])
|
115 |
-
|
116 |
-
#pdb.set_trace()
|
117 |
return results
|
118 |
|
119 |
##########################################
|
@@ -122,8 +121,6 @@ def crop_animal_detections(img_in,
|
|
122 |
likelihood_th):
|
123 |
|
124 |
## Extract animal crops
|
125 |
-
#print(yolo_results)
|
126 |
-
#pdb.set_trace()
|
127 |
list_labels_as_str = [i for i in yolo_results.names.values()] # ['animal', 'person', 'vehicle']
|
128 |
list_np_animal_crops = []
|
129 |
|
@@ -145,9 +142,7 @@ def crop_animal_detections(img_in,
|
|
145 |
|
146 |
pred_llk = det_array[j,4]
|
147 |
pred_label = det_array[j,5]
|
148 |
-
#pdb.set_trace()
|
149 |
# keep animal crops above threshold
|
150 |
-
#pdb.set_trace()
|
151 |
if (pred_label == list_labels_as_str.index('animal')) and \
|
152 |
(pred_llk >= likelihood_th):
|
153 |
area = (xmin_rd, ymin_rd, xmax_rd, ymax_rd)
|
@@ -162,7 +157,8 @@ def crop_animal_detections(img_in,
|
|
162 |
return list_np_animal_crops
|
163 |
|
164 |
def draw_rectangle_text(img,results,font_style='amiko',font_size=8, keypt_color="white",):
|
165 |
-
|
|
|
166 |
w, h = bbxyxy[2], bbxyxy[3]
|
167 |
shape = [(bbxyxy[0], bbxyxy[1]), (w , h)]
|
168 |
imgR = ImageDraw.Draw(img)
|
@@ -185,22 +181,24 @@ def predict_dlc(list_np_crops,
|
|
185 |
kpts_likelihood_th,
|
186 |
DLCmodel,
|
187 |
dlc_proc):
|
188 |
-
|
189 |
# run dlc thru list of crops
|
190 |
dlc_live = DLCLive(DLCmodel, processor=dlc_proc)
|
191 |
dlc_live.init_inference(list_np_crops[0])
|
192 |
|
193 |
list_kpts_per_crop = []
|
|
|
194 |
np_aux = np.empty((1,3)) # can I avoid hardcoding here?
|
195 |
for crop in list_np_crops:
|
196 |
# scale crop here?
|
197 |
keypts_xyp = dlc_live.get_pose(crop) # third column is llk!
|
198 |
# set kpts below threhsold to nan
|
199 |
-
|
200 |
#pdb.set_trace()
|
201 |
keypts_xyp[keypts_xyp[:,-1] < kpts_likelihood_th,:] = np_aux.fill(np.nan)
|
202 |
# add kpts of this crop to list
|
203 |
list_kpts_per_crop.append(keypts_xyp)
|
|
|
204 |
#return confidence here
|
205 |
return list_kpts_per_crop
|
206 |
|
@@ -230,6 +228,7 @@ def predict_pipeline(img_input,
|
|
230 |
# pose_cfg_dict['all_joints'] is a list of one-element lists,
|
231 |
with open(pose_cfg_path, "r") as stream:
|
232 |
pose_cfg_dict = yaml.safe_load(stream)
|
|
|
233 |
map_label_id_to_str = dict([(k,v) for k,v in zip([el[0] for el in pose_cfg_dict['all_joints']],
|
234 |
pose_cfg_dict['all_joints_names'])])
|
235 |
|
@@ -275,19 +274,10 @@ def predict_pipeline(img_input,
|
|
275 |
kpts_likelihood_th,
|
276 |
path_to_DLCmodel,
|
277 |
dlc_proc)
|
278 |
-
|
279 |
-
# Produce final image
|
280 |
-
# img_background = Image.fromarray(md_results.ims[0]) # img_input or Image.fromarray(md_results.imgs[0])?
|
281 |
-
# Image.fromarray(md_results.imgs[0]) --> (640, 479)
|
282 |
-
# img_input.size ---> (259, 194)
|
283 |
|
284 |
-
img_background = img_input.resize((md_results.ims[0].shape[1],
|
285 |
-
|
286 |
-
|
287 |
-
|
288 |
-
# resize image to match megadetector output
|
289 |
-
# g = (640 / max(img_background.size)) # gain
|
290 |
-
# img_background = img_background.resize((int(x * g) for x in img_background.size), Image.ANTIALIAS) # resize
|
291 |
for ic, (np_crop, kpts_crop) in enumerate(zip(list_crops,
|
292 |
list_kpts_per_crop)):
|
293 |
|
@@ -304,17 +294,23 @@ def predict_pipeline(img_input,
|
|
304 |
marker_size=marker_size)
|
305 |
|
306 |
## Paste crop in original image https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.paste
|
307 |
-
img_background.paste(img_crop,
|
308 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
309 |
|
310 |
-
|
311 |
-
|
312 |
-
|
|
|
|
|
313 |
|
314 |
#############################################
|
315 |
-
# %%
|
316 |
# User interface: inputs
|
317 |
-
|
318 |
# Input image
|
319 |
gr_image_input = gr.inputs.Image(type="pil", label="Input Image")
|
320 |
|
@@ -334,7 +330,7 @@ gr_dlc_only_checkbox = gr.inputs.Checkbox(False,
|
|
334 |
gr_str_labels_checkbox = gr.inputs.Checkbox(True,
|
335 |
label='Show bodypart labels?')
|
336 |
|
337 |
-
gr_slider_conf_bboxes = gr.inputs.Slider(0,1,.
|
338 |
label='Set confidence threshold for animal detections')
|
339 |
gr_slider_conf_keypoints = gr.inputs.Slider(0,1,.05,0,
|
340 |
label='Set confidence threshold for keypoints')
|
@@ -368,7 +364,8 @@ inputs = [gr_image_input,
|
|
368 |
# %%
|
369 |
# User interface: outputs
|
370 |
gr_image_output = gr.outputs.Image(type="pil", label="Output Image")
|
371 |
-
|
|
|
372 |
|
373 |
##############################################
|
374 |
# User interace: description
|
@@ -382,7 +379,7 @@ gr_description = "Contributed by Sofia Minano, Neslihan Wittek, Nirel Kadzo, Vic
|
|
382 |
# article = "<p style='text-align: center'>This app makes predictions using a YOLOv5x6 model that was trained to detect animals, humans, and vehicles in camera trap images; find out more about the project on <a href='https://github.com/microsoft/CameraTraps'>GitHub</a>. This app was built by Henry Lydecker but really depends on code and models developed by <a href='http://ecologize.org/'>Ecologize</a> and <a href='http://aka.ms/aiforearth'>Microsoft AI for Earth</a>. Find out more about the YOLO model from the original creator, <a href='https://pjreddie.com/darknet/yolo/'>Joseph Redmon</a>. YOLOv5 is a family of compound-scaled object detection models trained on the COCO dataset and developed by Ultralytics, and includes simple functionality for Test Time Augmentation (TTA), model ensembling, hyperparameter evolution, and export to ONNX, CoreML and TFLite. <a href='https://github.com/ultralytics/yolov5'>Source code</a> | <a href='https://pytorch.org/hub/ultralytics_yolov5'>PyTorch Hub</a></p>"
|
383 |
|
384 |
examples = [['example/monkey_full.jpg', 'md_v5a','full_macaque', False, True, 0.5, 0.3, 'amiko', 9, 'blue', 3],
|
385 |
-
['example/dog.jpeg', 'md_v5a', 'full_dog', False, True, 0.5, 0.
|
386 |
['example/cat.jpg', 'md_v5a', 'full_cat', False, True, 0.5, 0.05, 'amiko', 9, 'purple', 3]]
|
387 |
|
388 |
################################################
|
|
|
19 |
import os
|
20 |
import yaml
|
21 |
from model.models import DownloadModel
|
22 |
+
from save_results import save_results
|
23 |
import pdb
|
24 |
|
25 |
#########################################
|
|
|
112 |
|
113 |
## detect objects
|
114 |
results = MD_model(im) # inference # vars(results).keys()= dict_keys(['imgs', 'pred', 'names', 'files', 'times', 'xyxy', 'xywh', 'xyxyn', 'xywhn', 'n', 't', 's'])
|
115 |
+
|
|
|
116 |
return results
|
117 |
|
118 |
##########################################
|
|
|
121 |
likelihood_th):
|
122 |
|
123 |
## Extract animal crops
|
|
|
|
|
124 |
list_labels_as_str = [i for i in yolo_results.names.values()] # ['animal', 'person', 'vehicle']
|
125 |
list_np_animal_crops = []
|
126 |
|
|
|
142 |
|
143 |
pred_llk = det_array[j,4]
|
144 |
pred_label = det_array[j,5]
|
|
|
145 |
# keep animal crops above threshold
|
|
|
146 |
if (pred_label == list_labels_as_str.index('animal')) and \
|
147 |
(pred_llk >= likelihood_th):
|
148 |
area = (xmin_rd, ymin_rd, xmax_rd, ymax_rd)
|
|
|
157 |
return list_np_animal_crops
|
158 |
|
159 |
def draw_rectangle_text(img,results,font_style='amiko',font_size=8, keypt_color="white",):
|
160 |
+
#pdb.set_trace()
|
161 |
+
bbxyxy = results
|
162 |
w, h = bbxyxy[2], bbxyxy[3]
|
163 |
shape = [(bbxyxy[0], bbxyxy[1]), (w , h)]
|
164 |
imgR = ImageDraw.Draw(img)
|
|
|
181 |
kpts_likelihood_th,
|
182 |
DLCmodel,
|
183 |
dlc_proc):
|
184 |
+
|
185 |
# run dlc thru list of crops
|
186 |
dlc_live = DLCLive(DLCmodel, processor=dlc_proc)
|
187 |
dlc_live.init_inference(list_np_crops[0])
|
188 |
|
189 |
list_kpts_per_crop = []
|
190 |
+
all_kypts = []
|
191 |
np_aux = np.empty((1,3)) # can I avoid hardcoding here?
|
192 |
for crop in list_np_crops:
|
193 |
# scale crop here?
|
194 |
keypts_xyp = dlc_live.get_pose(crop) # third column is llk!
|
195 |
# set kpts below threhsold to nan
|
196 |
+
|
197 |
#pdb.set_trace()
|
198 |
keypts_xyp[keypts_xyp[:,-1] < kpts_likelihood_th,:] = np_aux.fill(np.nan)
|
199 |
# add kpts of this crop to list
|
200 |
list_kpts_per_crop.append(keypts_xyp)
|
201 |
+
all_kypts.append(keypts_xyp)
|
202 |
#return confidence here
|
203 |
return list_kpts_per_crop
|
204 |
|
|
|
228 |
# pose_cfg_dict['all_joints'] is a list of one-element lists,
|
229 |
with open(pose_cfg_path, "r") as stream:
|
230 |
pose_cfg_dict = yaml.safe_load(stream)
|
231 |
+
|
232 |
map_label_id_to_str = dict([(k,v) for k,v in zip([el[0] for el in pose_cfg_dict['all_joints']],
|
233 |
pose_cfg_dict['all_joints_names'])])
|
234 |
|
|
|
274 |
kpts_likelihood_th,
|
275 |
path_to_DLCmodel,
|
276 |
dlc_proc)
|
|
|
|
|
|
|
|
|
|
|
277 |
|
278 |
+
img_background = img_input.resize((md_results.ims[0].shape[1],md_results.ims[0].shape[0]))
|
279 |
+
print('I have ' + str(len(list_crops)) + ' bounding box')
|
280 |
+
|
|
|
|
|
|
|
|
|
281 |
for ic, (np_crop, kpts_crop) in enumerate(zip(list_crops,
|
282 |
list_kpts_per_crop)):
|
283 |
|
|
|
294 |
marker_size=marker_size)
|
295 |
|
296 |
## Paste crop in original image https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.paste
|
297 |
+
img_background.paste(img_crop, box = tuple([int(t) for t in md_results.xyxy[0][ic,:2]]))
|
298 |
+
|
299 |
+
|
300 |
+
#set trh!! FIXME
|
301 |
+
bb_per_animal = md_results.xyxy[0].tolist()[ic]
|
302 |
+
pred = md_results.xyxy[0].tolist()[ic][4]
|
303 |
+
if bbox_likelihood_th < pred:
|
304 |
+
draw_rectangle_text(img_background, bb_per_animal ,font_style=font_style,font_size=font_size, keypt_color=keypt_color)
|
305 |
|
306 |
+
print(pred)
|
307 |
+
|
308 |
+
download_file = save_results(md_results,list_kpts_per_crop,map_label_id_to_str,bbox_likelihood_th)
|
309 |
+
|
310 |
+
return img_background, download_file
|
311 |
|
312 |
#############################################
|
|
|
313 |
# User interface: inputs
|
|
|
314 |
# Input image
|
315 |
gr_image_input = gr.inputs.Image(type="pil", label="Input Image")
|
316 |
|
|
|
330 |
gr_str_labels_checkbox = gr.inputs.Checkbox(True,
|
331 |
label='Show bodypart labels?')
|
332 |
|
333 |
+
gr_slider_conf_bboxes = gr.inputs.Slider(0,1,.02,0.8,
|
334 |
label='Set confidence threshold for animal detections')
|
335 |
gr_slider_conf_keypoints = gr.inputs.Slider(0,1,.05,0,
|
336 |
label='Set confidence threshold for keypoints')
|
|
|
364 |
# %%
|
365 |
# User interface: outputs
|
366 |
gr_image_output = gr.outputs.Image(type="pil", label="Output Image")
|
367 |
+
out_smpl_npy_download = gr.File(label="Download JSON file")
|
368 |
+
outputs = [gr_image_output,out_smpl_npy_download]
|
369 |
|
370 |
##############################################
|
371 |
# User interace: description
|
|
|
379 |
# article = "<p style='text-align: center'>This app makes predictions using a YOLOv5x6 model that was trained to detect animals, humans, and vehicles in camera trap images; find out more about the project on <a href='https://github.com/microsoft/CameraTraps'>GitHub</a>. This app was built by Henry Lydecker but really depends on code and models developed by <a href='http://ecologize.org/'>Ecologize</a> and <a href='http://aka.ms/aiforearth'>Microsoft AI for Earth</a>. Find out more about the YOLO model from the original creator, <a href='https://pjreddie.com/darknet/yolo/'>Joseph Redmon</a>. YOLOv5 is a family of compound-scaled object detection models trained on the COCO dataset and developed by Ultralytics, and includes simple functionality for Test Time Augmentation (TTA), model ensembling, hyperparameter evolution, and export to ONNX, CoreML and TFLite. <a href='https://github.com/ultralytics/yolov5'>Source code</a> | <a href='https://pytorch.org/hub/ultralytics_yolov5'>PyTorch Hub</a></p>"
|
380 |
|
381 |
examples = [['example/monkey_full.jpg', 'md_v5a','full_macaque', False, True, 0.5, 0.3, 'amiko', 9, 'blue', 3],
|
382 |
+
['example/dog.jpeg', 'md_v5a', 'full_dog', False, True, 0.5, 0.00, 'amiko',9, 'yellow', 3],
|
383 |
['example/cat.jpg', 'md_v5a', 'full_cat', False, True, 0.5, 0.05, 'amiko', 9, 'purple', 3]]
|
384 |
|
385 |
################################################
|
save_results.py
ADDED
@@ -0,0 +1,59 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import json
|
2 |
+
import numpy as np
|
3 |
+
import pdb
|
4 |
+
|
5 |
+
dict_pred = {0: 'animal', 1: 'person', 2: 'vehicle'}
|
6 |
+
|
7 |
+
|
8 |
+
def save_results(md_results, dlc_outputs,map_label_id_to_str,thr,output_file = 'dowload_predictions.json'):
|
9 |
+
|
10 |
+
"""
|
11 |
+
|
12 |
+
write json
|
13 |
+
|
14 |
+
"""
|
15 |
+
info = {}
|
16 |
+
## info megaDetector
|
17 |
+
info['file']= md_results.files[0]
|
18 |
+
number_bb = len(md_results.xyxy[0].tolist())
|
19 |
+
info['number_of_bb'] = number_bb
|
20 |
+
number_bb_thr = len(dlc_outputs)
|
21 |
+
labels = [n for n in map_label_id_to_str.values()]
|
22 |
+
#pdb.set_trace()
|
23 |
+
new_index = []
|
24 |
+
for i in range(number_bb):
|
25 |
+
corner_x1,corner_y1,corner_x2,corner_y2,confidence, _ = md_results.xyxy[0].tolist()[i]
|
26 |
+
|
27 |
+
if confidence > thr:
|
28 |
+
new_index.append(i)
|
29 |
+
|
30 |
+
|
31 |
+
for i in range(number_bb_thr):
|
32 |
+
aux={}
|
33 |
+
corner_x1,corner_y1,corner_x2,corner_y2,confidence, _ = md_results.xyxy[0].tolist()[new_index[i]]
|
34 |
+
aux['corner_1'] = (corner_x1,corner_y1)
|
35 |
+
aux['corner_2'] = (corner_x2,corner_y2)
|
36 |
+
aux['predict MD'] = md_results.names[0]
|
37 |
+
aux['confidence MD'] = confidence
|
38 |
+
|
39 |
+
## info dlc
|
40 |
+
kypts = []
|
41 |
+
pdb.set_trace()
|
42 |
+
for s in dlc_outputs[i]:
|
43 |
+
#print(s)
|
44 |
+
aux1 = []
|
45 |
+
for j in s:
|
46 |
+
aux1.append(float(j))
|
47 |
+
|
48 |
+
kypts.append(aux1)
|
49 |
+
pdb.set_trace()
|
50 |
+
aux['dlc_pred'] = dict(zip(labels,kypts))
|
51 |
+
info['bb_' + str(new_index[i]) ]=aux
|
52 |
+
|
53 |
+
|
54 |
+
with open(output_file, 'w') as f:
|
55 |
+
json.dump(info, f, indent=1)
|
56 |
+
print('Output file saved at {}'.format(output_file))
|
57 |
+
|
58 |
+
return output_file
|
59 |
+
|