Spaces:
Running
Running
narugo1992
commited on
Commit
•
ebc32f0
1
Parent(s):
2dae70b
dev(narugo): add person detection
Browse files
app.py
CHANGED
@@ -3,6 +3,7 @@ import os
|
|
3 |
import gradio as gr
|
4 |
|
5 |
from face import _FACE_MODELS, _DEFAULT_FACE_MODEL, _gr_detect_faces
|
|
|
6 |
|
7 |
if __name__ == '__main__':
|
8 |
with gr.Blocks() as demo:
|
@@ -31,4 +32,28 @@ if __name__ == '__main__':
|
|
31 |
outputs=[gr_face_output_image],
|
32 |
)
|
33 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
34 |
demo.queue(os.cpu_count()).launch()
|
|
|
3 |
import gradio as gr
|
4 |
|
5 |
from face import _FACE_MODELS, _DEFAULT_FACE_MODEL, _gr_detect_faces
|
6 |
+
from person import _PERSON_MODELS, _DEFAULT_PERSON_MODEL, _gr_detect_person
|
7 |
|
8 |
if __name__ == '__main__':
|
9 |
with gr.Blocks() as demo:
|
|
|
32 |
outputs=[gr_face_output_image],
|
33 |
)
|
34 |
|
35 |
+
with gr.Tab('Person Detection'):
|
36 |
+
with gr.Row():
|
37 |
+
with gr.Column():
|
38 |
+
gr_person_input_image = gr.Image(type='pil', label='Original Image')
|
39 |
+
gr_person_model = gr.Dropdown(_PERSON_MODELS, value=_DEFAULT_PERSON_MODEL, label='Model')
|
40 |
+
gr_person_infer_size = gr.Slider(480, 1600, value=1216, step=32, label='Max Infer Size')
|
41 |
+
with gr.Row():
|
42 |
+
gr_person_iou_threshold = gr.Slider(0.0, 1.0, 0.5, label='IOU Threshold')
|
43 |
+
gr_person_score_threshold = gr.Slider(0.0, 1.0, 0.3, label='Score Threshold')
|
44 |
+
|
45 |
+
gr_person_submit = gr.Button(value='Submit', variant='primary')
|
46 |
+
|
47 |
+
with gr.Column():
|
48 |
+
gr_person_output_image = gr.Image(type='pil', label="Labeled")
|
49 |
+
|
50 |
+
gr_person_submit.click(
|
51 |
+
_gr_detect_person,
|
52 |
+
inputs=[
|
53 |
+
gr_person_input_image, gr_person_model,
|
54 |
+
gr_person_infer_size, gr_person_score_threshold, gr_person_iou_threshold,
|
55 |
+
],
|
56 |
+
outputs=[gr_person_output_image],
|
57 |
+
)
|
58 |
+
|
59 |
demo.queue(os.cpu_count()).launch()
|
person.py
ADDED
@@ -0,0 +1,38 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from functools import lru_cache
|
2 |
+
|
3 |
+
from huggingface_hub import hf_hub_download
|
4 |
+
from imgutils.data import ImageTyping, load_image, rgb_encode
|
5 |
+
|
6 |
+
from onnx_ import _open_onnx_model
|
7 |
+
from plot import plot_detection
|
8 |
+
from yolo_ import _image_preprocess, _data_simple_postprocess
|
9 |
+
|
10 |
+
_PERSON_MODELS = [
|
11 |
+
'person_detect_best_s.onnx',
|
12 |
+
]
|
13 |
+
_DEFAULT_PERSON_MODEL = _PERSON_MODELS[0]
|
14 |
+
|
15 |
+
|
16 |
+
@lru_cache()
|
17 |
+
def _open_person_detect_model(model_name):
|
18 |
+
return _open_onnx_model(hf_hub_download(
|
19 |
+
'deepghs/imgutils-models',
|
20 |
+
f'person_detect/{model_name}'
|
21 |
+
))
|
22 |
+
|
23 |
+
|
24 |
+
def detect_person(image: ImageTyping, model_name: str, max_infer_size=1216,
|
25 |
+
conf_threshold: float = 0.25, iou_threshold: float = 0.7):
|
26 |
+
image = load_image(image, mode='RGB')
|
27 |
+
new_image, old_size, new_size = _image_preprocess(image, max_infer_size)
|
28 |
+
|
29 |
+
data = rgb_encode(new_image)[None, ...]
|
30 |
+
output, = _open_person_detect_model(model_name).run(['output0'], {'images': data})
|
31 |
+
return _data_simple_postprocess(output[0], conf_threshold, iou_threshold, old_size, new_size)
|
32 |
+
|
33 |
+
|
34 |
+
def _gr_detect_person(image: ImageTyping, model_name: str, max_infer_size=1216,
|
35 |
+
conf_threshold: float = 0.25, iou_threshold: float = 0.7):
|
36 |
+
ret = detect_person(image, model_name, max_infer_size, conf_threshold, iou_threshold)
|
37 |
+
detections = [(box, 0, score) for box, score in ret]
|
38 |
+
return plot_detection(image, detections, ['person'])
|
yolo_.py
CHANGED
@@ -92,6 +92,9 @@ def _data_simple_postprocess(output, conf_threshold, iou_threshold, old_size, ne
|
|
92 |
scores = output[4, :]
|
93 |
records = sorted(zip(boxes, scores), key=lambda x: -x[1])
|
94 |
|
|
|
|
|
|
|
95 |
boxes = _yolo_xywh2xyxy(np.stack([bx for bx, _ in records]))
|
96 |
scores = np.stack([score for _, score in records])
|
97 |
idx = _yolo_nms(boxes, scores, thresh=iou_threshold)
|
|
|
92 |
scores = output[4, :]
|
93 |
records = sorted(zip(boxes, scores), key=lambda x: -x[1])
|
94 |
|
95 |
+
if not records:
|
96 |
+
return []
|
97 |
+
|
98 |
boxes = _yolo_xywh2xyxy(np.stack([bx for bx, _ in records]))
|
99 |
scores = np.stack([score for _, score in records])
|
100 |
idx = _yolo_nms(boxes, scores, thresh=iou_threshold)
|