narugo1992 commited on
Commit
528f68e
1 Parent(s): 59c1c71

dev(narugo): add halfbody, ci skip

Browse files
Files changed (2) hide show
  1. app.py +25 -0
  2. halfbody.py +43 -0
app.py CHANGED
@@ -4,6 +4,7 @@ import gradio as gr
4
 
5
  from censor import _CENSOR_MODELS, _DEFAULT_CENSOR_MODEL, _gr_detect_censors
6
  from face import _FACE_MODELS, _DEFAULT_FACE_MODEL, _gr_detect_faces
 
7
  from hand import _gr_detect_hands, _HAND_MODELS, _DEFAULT_HAND_MODEL
8
  from head import _gr_detect_heads, _HEAD_MODELS, _DEFAULT_HEAD_MODEL
9
  from manbits import _MANBIT_MODELS, _DEFAULT_MANBIT_MODEL, _gr_detect_manbits
@@ -84,6 +85,30 @@ if __name__ == '__main__':
84
  outputs=[gr_person_output_image],
85
  )
86
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
87
  with gr.Tab('Hand Detection'):
88
  with gr.Row():
89
  with gr.Column():
 
4
 
5
  from censor import _CENSOR_MODELS, _DEFAULT_CENSOR_MODEL, _gr_detect_censors
6
  from face import _FACE_MODELS, _DEFAULT_FACE_MODEL, _gr_detect_faces
7
+ from halfbody import _HALFBODY_MODELS, _DEFAULT_HALFBODY_MODEL, _gr_detect_halfbodies
8
  from hand import _gr_detect_hands, _HAND_MODELS, _DEFAULT_HAND_MODEL
9
  from head import _gr_detect_heads, _HEAD_MODELS, _DEFAULT_HEAD_MODEL
10
  from manbits import _MANBIT_MODELS, _DEFAULT_MANBIT_MODEL, _gr_detect_manbits
 
85
  outputs=[gr_person_output_image],
86
  )
87
 
88
+ with gr.Tab('Half Body Detection'):
89
+ with gr.Row():
90
+ with gr.Column():
91
+ gr_halfbody_input_image = gr.Image(type='pil', label='Original Image')
92
+ gr_halfbody_model = gr.Dropdown(_HALFBODY_MODELS, value=_DEFAULT_HALFBODY_MODEL, label='Model')
93
+ gr_halfbody_infer_size = gr.Slider(480, 960, value=640, step=32, label='Max Infer Size')
94
+ with gr.Row():
95
+ gr_halfbody_iou_threshold = gr.Slider(0.0, 1.0, 0.7, label='IOU Threshold')
96
+ gr_halfbody_score_threshold = gr.Slider(0.0, 1.0, 0.25, label='Score Threshold')
97
+
98
+ gr_halfbody_submit = gr.Button(value='Submit', variant='primary')
99
+
100
+ with gr.Column():
101
+ gr_halfbody_output_image = gr.Image(type='pil', label="Labeled")
102
+
103
+ gr_halfbody_submit.click(
104
+ _gr_detect_halfbodies,
105
+ inputs=[
106
+ gr_halfbody_input_image, gr_halfbody_model,
107
+ gr_halfbody_infer_size, gr_halfbody_score_threshold, gr_halfbody_iou_threshold,
108
+ ],
109
+ outputs=[gr_halfbody_output_image],
110
+ )
111
+
112
  with gr.Tab('Hand Detection'):
113
  with gr.Row():
114
  with gr.Column():
halfbody.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from functools import lru_cache
2
+ from typing import List, Tuple
3
+
4
+ from huggingface_hub import hf_hub_download
5
+ from imgutils.data import ImageTyping, load_image, rgb_encode
6
+
7
+ from onnx_ import _open_onnx_model
8
+ from plot import detection_visualize
9
+ from yolo_ import _image_preprocess, _data_postprocess
10
+
11
+ _HALFBODY_MODELS = [
12
+ 'halfbody_detect_v0.3_s',
13
+ 'halfbody_detect_v0.2_s',
14
+ ]
15
+ _DEFAULT_HALFBODY_MODEL = _HALFBODY_MODELS[0]
16
+
17
+
18
+ @lru_cache()
19
+ def _open_halfbody_detect_model(model_name):
20
+ return _open_onnx_model(hf_hub_download(
21
+ f'deepghs/anime_halfbody_detection',
22
+ f'{model_name}/model.onnx'
23
+ ))
24
+
25
+
26
+ _LABELS = ['haldbody']
27
+
28
+
29
+ def detect_halfbodies(image: ImageTyping, model_name: str, max_infer_size=640,
30
+ conf_threshold: float = 0.25, iou_threshold: float = 0.5) \
31
+ -> List[Tuple[Tuple[int, int, int, int], str, float]]:
32
+ image = load_image(image, mode='RGB')
33
+ new_image, old_size, new_size = _image_preprocess(image, max_infer_size)
34
+
35
+ data = rgb_encode(new_image)[None, ...]
36
+ output, = _open_halfbody_detect_model(model_name).run(['output0'], {'images': data})
37
+ return _data_postprocess(output[0], conf_threshold, iou_threshold, old_size, new_size, _LABELS)
38
+
39
+
40
+ def _gr_detect_halfbodies(image: ImageTyping, model_name: str, max_infer_size=640,
41
+ conf_threshold: float = 0.25, iou_threshold: float = 0.5):
42
+ ret = detect_halfbodies(image, model_name, max_infer_size, conf_threshold, iou_threshold)
43
+ return detection_visualize(image, ret, _LABELS)