dev(narugo): add person detection

app.py

@@ -3,6 +3,7 @@ import os
 import gradio as gr
 
 from face import _FACE_MODELS, _DEFAULT_FACE_MODEL, _gr_detect_faces
+from person import _PERSON_MODELS, _DEFAULT_PERSON_MODEL, _gr_detect_person
 
 if __name__ == '__main__':
     with gr.Blocks() as demo:
@@ -31,4 +32,28 @@ if __name__ == '__main__':
                         outputs=[gr_face_output_image],
                     )
 
+            with gr.Tab('Person Detection'):
+                with gr.Row():
+                    with gr.Column():
+                        gr_person_input_image = gr.Image(type='pil', label='Original Image')
+                        gr_person_model = gr.Dropdown(_PERSON_MODELS, value=_DEFAULT_PERSON_MODEL, label='Model')
+                        gr_person_infer_size = gr.Slider(480, 1600, value=1216, step=32, label='Max Infer Size')
+                        with gr.Row():
+                            gr_person_iou_threshold = gr.Slider(0.0, 1.0, 0.5, label='IOU Threshold')
+                            gr_person_score_threshold = gr.Slider(0.0, 1.0, 0.3, label='Score Threshold')
+
+                        gr_person_submit = gr.Button(value='Submit', variant='primary')
+
+                    with gr.Column():
+                        gr_person_output_image = gr.Image(type='pil', label="Labeled")
+
+                    gr_person_submit.click(
+                        _gr_detect_person,
+                        inputs=[
+                            gr_person_input_image, gr_person_model,
+                            gr_person_infer_size, gr_person_score_threshold, gr_person_iou_threshold,
+                        ],
+                        outputs=[gr_person_output_image],
+                    )
+
     demo.queue(os.cpu_count()).launch()

person.py

@@ -0,0 +1,38 @@
+from functools import lru_cache
+
+from huggingface_hub import hf_hub_download
+from imgutils.data import ImageTyping, load_image, rgb_encode
+
+from onnx_ import _open_onnx_model
+from plot import plot_detection
+from yolo_ import _image_preprocess, _data_simple_postprocess
+
+_PERSON_MODELS = [
+    'person_detect_best_s.onnx',
+]
+_DEFAULT_PERSON_MODEL = _PERSON_MODELS[0]
+
+
+@lru_cache()
+def _open_person_detect_model(model_name):
+    return _open_onnx_model(hf_hub_download(
+        'deepghs/imgutils-models',
+        f'person_detect/{model_name}'
+    ))
+
+
+def detect_person(image: ImageTyping, model_name: str, max_infer_size=1216,
+                  conf_threshold: float = 0.25, iou_threshold: float = 0.7):
+    image = load_image(image, mode='RGB')
+    new_image, old_size, new_size = _image_preprocess(image, max_infer_size)
+
+    data = rgb_encode(new_image)[None, ...]
+    output, = _open_person_detect_model(model_name).run(['output0'], {'images': data})
+    return _data_simple_postprocess(output[0], conf_threshold, iou_threshold, old_size, new_size)
+
+
+def _gr_detect_person(image: ImageTyping, model_name: str, max_infer_size=1216,
+                      conf_threshold: float = 0.25, iou_threshold: float = 0.7):
+    ret = detect_person(image, model_name, max_infer_size, conf_threshold, iou_threshold)
+    detections = [(box, 0, score) for box, score in ret]
+    return plot_detection(image, detections, ['person'])

yolo_.py

@@ -92,6 +92,9 @@ def _data_simple_postprocess(output, conf_threshold, iou_threshold, old_size, ne
     scores = output[4, :]
     records = sorted(zip(boxes, scores), key=lambda x: -x[1])
 
+    if not records:
+        return []
+
     boxes = _yolo_xywh2xyxy(np.stack([bx for bx, _ in records]))
     scores = np.stack([score for _, score in records])
     idx = _yolo_nms(boxes, scores, thresh=iou_threshold)