dev(narugo): upgrade this space

README.md

@@ -8,6 +8,14 @@ sdk_version: 4.44.0
 app_file: app.py
 pinned: false
 license: mit
+models:
+- deepghs/anime_face_detection
+- deepghs/anime_censor_detection
+- deepghs/anime_eye_detection
+- deepghs/anime_halfbody_detection
+- deepghs/anime_hand_detection
+- deepghs/anime_head_detection
+- deepghs/anime_person_detection
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -2,208 +2,39 @@ import os
 
 import gradio as gr
 
-from censor import _CENSOR_MODELS, _DEFAULT_CENSOR_MODEL, _gr_detect_censors
-from eye import _EYE_MODELS, _DEFAULT_EYE_MODEL, _gr_detect_eyes
-from face import _FACE_MODELS, _DEFAULT_FACE_MODEL, _gr_detect_faces
-from halfbody import _HALFBODY_MODELS, _DEFAULT_HALFBODY_MODEL, _gr_detect_halfbodies
-from hand import _gr_detect_hands, _HAND_MODELS, _DEFAULT_HAND_MODEL
-from head import _gr_detect_heads, _HEAD_MODELS, _DEFAULT_HEAD_MODEL
-from manbits import _MANBIT_MODELS, _DEFAULT_MANBIT_MODEL, _gr_detect_manbits
-from person import _PERSON_MODELS, _DEFAULT_PERSON_MODEL, _gr_detect_person
+from detection import EyesDetection, FaceDetection, HeadDetection, PersonDetection, HandDetection, CensorDetection, \
+    HalfBodyDetection
 
-if __name__ == '__main__':
-    with gr.Blocks() as demo:
-        with gr.Tabs():
-            with gr.Tab('Face Detection'):
-                with gr.Row():
-                    with gr.Column():
-                        gr_face_input_image = gr.Image(type='pil', label='Original Image')
-                        gr_face_model = gr.Dropdown(_FACE_MODELS, value=_DEFAULT_FACE_MODEL, label='Model')
-                        gr_face_infer_size = gr.Slider(480, 960, value=640, step=32, label='Max Infer Size')
-                        with gr.Row():
-                            gr_face_iou_threshold = gr.Slider(0.0, 1.0, 0.7, label='IOU Threshold')
-                            gr_face_score_threshold = gr.Slider(0.0, 1.0, 0.25, label='Score Threshold')
-
-                        gr_face_submit = gr.Button(value='Submit', variant='primary')
-
-                    with gr.Column():
-                        gr_face_output_image = gr.Image(type='pil', label="Labeled")
-
-                    gr_face_submit.click(
-                        _gr_detect_faces,
-                        inputs=[
-                            gr_face_input_image, gr_face_model,
-                            gr_face_infer_size, gr_face_score_threshold, gr_face_iou_threshold,
-                        ],
-                        outputs=[gr_face_output_image],
-                    )
-
-            with gr.Tab('Head Detection'):
-                with gr.Row():
-                    with gr.Column():
-                        gr_head_input_image = gr.Image(type='pil', label='Original Image')
-                        gr_head_model = gr.Dropdown(_HEAD_MODELS, value=_DEFAULT_HEAD_MODEL, label='Model')
-                        gr_head_infer_size = gr.Slider(480, 960, value=640, step=32, label='Max Infer Size')
-                        with gr.Row():
-                            gr_head_iou_threshold = gr.Slider(0.0, 1.0, 0.7, label='IOU Threshold')
-                            gr_head_score_threshold = gr.Slider(0.0, 1.0, 0.3, label='Score Threshold')
-
-                        gr_head_submit = gr.Button(value='Submit', variant='primary')
-
-                    with gr.Column():
-                        gr_head_output_image = gr.Image(type='pil', label="Labeled")
-
-                    gr_head_submit.click(
-                        _gr_detect_heads,
-                        inputs=[
-                            gr_head_input_image, gr_head_model,
-                            gr_head_infer_size, gr_head_score_threshold, gr_head_iou_threshold,
-                        ],
-                        outputs=[gr_head_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, 960, value=640, 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],
-                    )
-
-            with gr.Tab('Half Body Detection'):
-                with gr.Row():
-                    with gr.Column():
-                        gr_halfbody_input_image = gr.Image(type='pil', label='Original Image')
-                        gr_halfbody_model = gr.Dropdown(_HALFBODY_MODELS, value=_DEFAULT_HALFBODY_MODEL, label='Model')
-                        gr_halfbody_infer_size = gr.Slider(480, 960, value=640, step=32, label='Max Infer Size')
-                        with gr.Row():
-                            gr_halfbody_iou_threshold = gr.Slider(0.0, 1.0, 0.7, label='IOU Threshold')
-                            gr_halfbody_score_threshold = gr.Slider(0.0, 1.0, 0.25, label='Score Threshold')
-
-                        gr_halfbody_submit = gr.Button(value='Submit', variant='primary')
-
-                    with gr.Column():
-                        gr_halfbody_output_image = gr.Image(type='pil', label="Labeled")
-
-                    gr_halfbody_submit.click(
-                        _gr_detect_halfbodies,
-                        inputs=[
-                            gr_halfbody_input_image, gr_halfbody_model,
-                            gr_halfbody_infer_size, gr_halfbody_score_threshold, gr_halfbody_iou_threshold,
-                        ],
-                        outputs=[gr_halfbody_output_image],
-                    )
+_GLOBAL_CSS = """
+.limit-height {
+    max-height: 55vh;
+}
+"""
 
-            with gr.Tab('Eyes Detection'):
-                with gr.Row():
-                    with gr.Column():
-                        gr_eye_input_image = gr.Image(type='pil', label='Original Image')
-                        gr_eye_model = gr.Dropdown(_EYE_MODELS, value=_DEFAULT_EYE_MODEL, label='Model')
-                        gr_eye_infer_size = gr.Slider(480, 960, value=640, step=32, label='Max Infer Size')
-                        with gr.Row():
-                            gr_eye_iou_threshold = gr.Slider(0.0, 1.0, 0.3, label='IOU Threshold')
-                            gr_eye_score_threshold = gr.Slider(0.0, 1.0, 0.3, label='Score Threshold')
-
-                        gr_eye_submit = gr.Button(value='Submit', variant='primary')
-
-                    with gr.Column():
-                        gr_eye_output_image = gr.Image(type='pil', label="Labeled")
-
-                    gr_eye_submit.click(
-                        _gr_detect_eyes,
-                        inputs=[
-                            gr_eye_input_image, gr_eye_model,
-                            gr_eye_infer_size, gr_eye_score_threshold, gr_eye_iou_threshold,
-                        ],
-                        outputs=[gr_eye_output_image],
-                    )
-
-            with gr.Tab('Hand Detection'):
-                with gr.Row():
-                    with gr.Column():
-                        gr_hand_input_image = gr.Image(type='pil', label='Original Image')
-                        gr_hand_model = gr.Dropdown(_HAND_MODELS, value=_DEFAULT_HAND_MODEL, label='Model')
-                        gr_hand_infer_size = gr.Slider(480, 960, value=640, step=32, label='Max Infer Size')
-                        with gr.Row():
-                            gr_hand_iou_threshold = gr.Slider(0.0, 1.0, 0.7, label='IOU Threshold')
-                            gr_hand_score_threshold = gr.Slider(0.0, 1.0, 0.35, label='Score Threshold')
-
-                        gr_hand_submit = gr.Button(value='Submit', variant='primary')
-
-                    with gr.Column():
-                        gr_hand_output_image = gr.Image(type='pil', label="Labeled")
-
-                    gr_hand_submit.click(
-                        _gr_detect_hands,
-                        inputs=[
-                            gr_hand_input_image, gr_hand_model,
-                            gr_hand_infer_size, gr_hand_score_threshold, gr_hand_iou_threshold,
-                        ],
-                        outputs=[gr_hand_output_image],
-                    )
-
-            with gr.Tab('Censor Point Detection'):
-                with gr.Row():
-                    with gr.Column():
-                        gr_censor_input_image = gr.Image(type='pil', label='Original Image')
-                        gr_censor_model = gr.Dropdown(_CENSOR_MODELS, value=_DEFAULT_CENSOR_MODEL, label='Model')
-                        gr_censor_infer_size = gr.Slider(480, 960, value=640, step=32, label='Max Infer Size')
-                        with gr.Row():
-                            gr_censor_iou_threshold = gr.Slider(0.0, 1.0, 0.5, label='IOU Threshold')
-                            gr_censor_score_threshold = gr.Slider(0.0, 1.0, 0.25, label='Score Threshold')
-
-                        gr_censor_submit = gr.Button(value='Submit', variant='primary')
-
-                    with gr.Column():
-                        gr_censor_output_image = gr.Image(type='pil', label="Labeled")
-
-                    gr_censor_submit.click(
-                        _gr_detect_censors,
-                        inputs=[
-                            gr_censor_input_image, gr_censor_model,
-                            gr_censor_infer_size, gr_censor_score_threshold, gr_censor_iou_threshold,
-                        ],
-                        outputs=[gr_censor_output_image],
-                    )
-
-            with gr.Tab('Manbits Detection\n(Deprecated)'):
-                with gr.Row():
-                    with gr.Column():
-                        gr_manbit_input_image = gr.Image(type='pil', label='Original Image')
-                        gr_manbit_model = gr.Dropdown(_MANBIT_MODELS, value=_DEFAULT_MANBIT_MODEL, label='Model')
-                        gr_manbit_infer_size = gr.Slider(480, 960, value=640, step=32, label='Max Infer Size')
-                        with gr.Row():
-                            gr_manbit_iou_threshold = gr.Slider(0.0, 1.0, 0.7, label='IOU Threshold')
-                            gr_manbit_score_threshold = gr.Slider(0.0, 1.0, 0.25, label='Score Threshold')
-
-                        gr_manbit_submit = gr.Button(value='Submit', variant='primary')
-
-                    with gr.Column():
-                        gr_manbit_output_image = gr.Image(type='pil', label="Labeled")
-
-                    gr_manbit_submit.click(
-                        _gr_detect_manbits,
-                        inputs=[
-                            gr_manbit_input_image, gr_manbit_model,
-                            gr_manbit_infer_size, gr_manbit_score_threshold, gr_manbit_iou_threshold,
-                        ],
-                        outputs=[gr_manbit_output_image],
-                    )
+if __name__ == '__main__':
+    with gr.Blocks(css=_GLOBAL_CSS) as demo:
+        with gr.Row():
+            with gr.Column():
+                gr.HTML('<h2 style="text-align: center;">Object Detections For Anime</h2>')
+                gr.Markdown('This is the online demo for detection functions of '
+                            '[imgutils.detect](https://dghs-imgutils.deepghs.org/main/api_doc/detect/index.html). '
+                            'You can try them yourselves with `pip install dghs-imgutils`.')
+
+        with gr.Row():
+            with gr.Tabs():
+                with gr.Tab('Face Detection'):
+                    FaceDetection().make_ui()
+                with gr.Tab('Head Detection'):
+                    HeadDetection().make_ui()
+                with gr.Tab('Person Detection'):
+                    PersonDetection().make_ui()
+                with gr.Tab('Half Body Detection'):
+                    HalfBodyDetection().make_ui()
+                with gr.Tab('Eyes Detection'):
+                    EyesDetection().make_ui()
+                with gr.Tab('Hand Detection'):
+                    HandDetection().make_ui()
+                with gr.Tab('Censor Point Detection'):
+                    CensorDetection().make_ui()
 
     demo.queue(os.cpu_count()).launch()

app2.py

@@ -1,34 +0,0 @@
-import os
-
-import gradio as gr
-
-from detection import EyesDetection, FaceDetection, HeadDetection, PersonDetection, HandDetection, CensorDetection, \
-    HalfBodyDetection
-
-_GLOBAL_CSS = """
-.limit-height {
-    max-height: 55vh;
-}
-"""
-
-if __name__ == '__main__':
-    with gr.Blocks(css=_GLOBAL_CSS) as demo:
-        with gr.Tabs():
-            with gr.Tab('Face Detection'):
-                FaceDetection().make_ui()
-            with gr.Tab('Head Detection'):
-                HeadDetection().make_ui()
-            with gr.Tab('Person Detection'):
-                PersonDetection().make_ui()
-            with gr.Tab('Half Body Detection'):
-                HalfBodyDetection().make_ui()
-            with gr.Tab('Eyes Detection'):
-                EyesDetection().make_ui()
-            with gr.Tab('Hand Detection'):
-                HandDetection().make_ui()
-            with gr.Tab('Censor Point Detection'):
-                CensorDetection().make_ui()
-            with gr.Tab('Manbits Detection\n(Deprecated)'):
-                pass
-
-    demo.queue(os.cpu_count()).launch()

censor.py

@@ -1,47 +0,0 @@
-from functools import lru_cache
-from typing import List, Tuple
-
-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 detection_visualize
-from yolo_ import _image_preprocess, _data_postprocess
-
-_CENSOR_MODELS = [
-    'censor_detect_v1.0_s',
-    'censor_detect_v1.0_n',
-    'censor_detect_v0.10_s',
-    'censor_detect_v0.9_s',
-    'censor_detect_v0.8_s',
-    'censor_detect_v0.7_s',
-]
-_DEFAULT_CENSOR_MODEL = _CENSOR_MODELS[0]
-
-
-@lru_cache()
-def _open_censor_detect_model(model_name):
-    return _open_onnx_model(hf_hub_download(
-        f'deepghs/anime_censor_detection',
-        f'{model_name}/model.onnx'
-    ))
-
-
-_LABELS = ['nipple_f', 'penis', 'pussy']
-
-
-def detect_censors(image: ImageTyping, model_name: str, max_infer_size=640,
-                   conf_threshold: float = 0.25, iou_threshold: float = 0.5) \
-        -> List[Tuple[Tuple[int, int, int, int], str, float]]:
-    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_censor_detect_model(model_name).run(['output0'], {'images': data})
-    return _data_postprocess(output[0], conf_threshold, iou_threshold, old_size, new_size, _LABELS)
-
-
-def _gr_detect_censors(image: ImageTyping, model_name: str, max_infer_size=640,
-                       conf_threshold: float = 0.25, iou_threshold: float = 0.5):
-    ret = detect_censors(image, model_name, max_infer_size, conf_threshold, iou_threshold)
-    return detection_visualize(image, ret, _LABELS)

detection/halfbody.py

@@ -0,0 +1,33 @@
+import re
+from typing import List, Tuple
+
+from imgutils.data import ImageTyping
+from imgutils.detect.halfbody import detect_halfbody, _LABELS
+
+from .base import DeepGHSObjectDetection
+
+
+def _parse_model_name(model_name: str):
+    matching = re.fullmatch(r'^halfbody_detect_(?P<version>[\s\S]+?)_(?P<level>[\s\S]+?)$', model_name)
+    return matching.group('version'), matching.group('level')
+
+
+class HalfBodyDetection(DeepGHSObjectDetection):
+    def __init__(self):
+        DeepGHSObjectDetection.__init__(self, repo_id='deepghs/anime_halfbody_detection')
+
+    def _get_default_model(self) -> str:
+        return 'halfbody_detect_v1.0_s'
+
+    def _get_default_iou_and_score(self, model_name: str) -> Tuple[float, float]:
+        return 0.7, 0.5
+
+    def _get_labels(self, model_name: str) -> List[str]:
+        return _LABELS
+
+    def detect(self, image: ImageTyping, model_name: str,
+               iou_threshold: float = 0.7, score_threshold: float = 0.25) \
+            -> List[Tuple[Tuple[float, float, float, float], str, float]]:
+        version, level = _parse_model_name(model_name)
+        return detect_halfbody(image, level=level, version=version,
+                               conf_threshold=score_threshold, iou_threshold=iou_threshold)

eye.py

@@ -1,50 +0,0 @@
-from functools import lru_cache
-from typing import List, Tuple
-
-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 detection_visualize
-from yolo_ import _image_preprocess, _data_postprocess
-
-_EYE_MODELS = [
-    'eye_detect_v1.0_s',
-    'eye_detect_v1.0_n',
-    'eye_detect_v0.8_s',
-    'eye_detect_v0.7_s',
-    'eye_detect_v0.6_s',
-    'eye_detect_v0.5_s',
-    'eye_detect_v0.4_s',
-    'eye_detect_v0.3_s',
-    'eye_detect_v0.2_s',
-]
-_DEFAULT_EYE_MODEL = _EYE_MODELS[0]
-
-
-@lru_cache()
-def _open_eye_detect_model(model_name):
-    return _open_onnx_model(hf_hub_download(
-        f'deepghs/anime_eye_detection',
-        f'{model_name}/model.onnx'
-    ))
-
-
-_LABELS = ['eye']
-
-
-def detect_eyes(image: ImageTyping, model_name: str, max_infer_size=640,
-                conf_threshold: float = 0.3, iou_threshold: float = 0.3) \
-        -> List[Tuple[Tuple[int, int, int, int], str, float]]:
-    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_eye_detect_model(model_name).run(['output0'], {'images': data})
-    return _data_postprocess(output[0], conf_threshold, iou_threshold, old_size, new_size, _LABELS)
-
-
-def _gr_detect_eyes(image: ImageTyping, model_name: str, max_infer_size=640,
-                    conf_threshold: float = 0.3, iou_threshold: float = 0.3):
-    ret = detect_eyes(image, model_name, max_infer_size, conf_threshold, iou_threshold)
-    return detection_visualize(image, ret, _LABELS)

face.py

@@ -1,52 +0,0 @@
-from functools import lru_cache
-from typing import List, Tuple
-
-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 detection_visualize
-from yolo_ import _image_preprocess, _data_postprocess
-
-_FACE_MODELS = [
-    'face_detect_v1.4_s',
-    'face_detect_v1.4_n',
-    'face_detect_v1.3_s',
-    'face_detect_v1.3_n',
-    'face_detect_v1.2_s',
-    'face_detect_v1.1_s',
-    'face_detect_v1.1_n',
-    'face_detect_v1_s',
-    'face_detect_v1_n',
-    'face_detect_v0_s',
-    'face_detect_v0_n',
-]
-_DEFAULT_FACE_MODEL = _FACE_MODELS[0]
-
-
-@lru_cache()
-def _open_face_detect_model(model_name):
-    return _open_onnx_model(hf_hub_download(
-        f'deepghs/anime_face_detection',
-        f'{model_name}/model.onnx',
-    ))
-
-
-_LABELS = ['face']
-
-
-def detect_faces(image: ImageTyping, model_name: str, max_infer_size=640,
-                 conf_threshold: float = 0.25, iou_threshold: float = 0.7) \
-        -> List[Tuple[Tuple[int, int, int, int], str, float]]:
-    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_face_detect_model(model_name).run(['output0'], {'images': data})
-    return _data_postprocess(output[0], conf_threshold, iou_threshold, old_size, new_size, _LABELS)
-
-
-def _gr_detect_faces(image: ImageTyping, model_name: str, max_infer_size=640,
-                     conf_threshold: float = 0.25, iou_threshold: float = 0.7):
-    ret = detect_faces(image, model_name, max_infer_size, conf_threshold, iou_threshold)
-    return detection_visualize(image, ret, _LABELS)

halfbody.py

@@ -1,44 +0,0 @@
-from functools import lru_cache
-from typing import List, Tuple
-
-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 detection_visualize
-from yolo_ import _image_preprocess, _data_postprocess
-
-_HALFBODY_MODELS = [
-    'halfbody_detect_v0.4_s',
-    'halfbody_detect_v0.3_s',
-    'halfbody_detect_v0.2_s',
-]
-_DEFAULT_HALFBODY_MODEL = _HALFBODY_MODELS[0]
-
-
-@lru_cache()
-def _open_halfbody_detect_model(model_name):
-    return _open_onnx_model(hf_hub_download(
-        f'deepghs/anime_halfbody_detection',
-        f'{model_name}/model.onnx'
-    ))
-
-
-_LABELS = ['haldbody']
-
-
-def detect_halfbodies(image: ImageTyping, model_name: str, max_infer_size=640,
-                      conf_threshold: float = 0.25, iou_threshold: float = 0.5) \
-        -> List[Tuple[Tuple[int, int, int, int], str, float]]:
-    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_halfbody_detect_model(model_name).run(['output0'], {'images': data})
-    return _data_postprocess(output[0], conf_threshold, iou_threshold, old_size, new_size, _LABELS)
-
-
-def _gr_detect_halfbodies(image: ImageTyping, model_name: str, max_infer_size=640,
-                          conf_threshold: float = 0.25, iou_threshold: float = 0.5):
-    ret = detect_halfbodies(image, model_name, max_infer_size, conf_threshold, iou_threshold)
-    return detection_visualize(image, ret, _LABELS)

hand.py

@@ -1,49 +0,0 @@
-from functools import lru_cache
-from typing import List, Tuple
-
-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 detection_visualize
-from yolo_ import _image_preprocess, _data_postprocess
-
-_HAND_MODELS = [
-    'hand_detect_v0.7_s',
-    'hand_detect_v0.6_s',
-    'hand_detect_v0.5_s',
-    'hand_detect_v0.4_s',
-    'hand_detect_v0.3_s',
-    'hand_detect_v0.2_s',
-    'hand_detect_v0.1_s',
-    'hand_detect_v0.1_n',
-]
-_DEFAULT_HAND_MODEL = _HAND_MODELS[0]
-
-
-@lru_cache()
-def _open_hand_detect_model(model_name):
-    return _open_onnx_model(hf_hub_download(
-        f'deepghs/anime_hand_detection',
-        f'{model_name}/model.onnx'
-    ))
-
-
-_LABELS = ['hand']
-
-
-def detect_hands(image: ImageTyping, model_name: str, max_infer_size=640,
-                 conf_threshold: float = 0.35, iou_threshold: float = 0.7) \
-        -> List[Tuple[Tuple[int, int, int, int], str, float]]:
-    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_hand_detect_model(model_name).run(['output0'], {'images': data})
-    return _data_postprocess(output[0], conf_threshold, iou_threshold, old_size, new_size, _LABELS)
-
-
-def _gr_detect_hands(image: ImageTyping, model_name: str, max_infer_size=640,
-                     conf_threshold: float = 0.35, iou_threshold: float = 0.7):
-    ret = detect_hands(image, model_name, max_infer_size, conf_threshold, iou_threshold)
-    return detection_visualize(image, ret, _LABELS)

head.py

@@ -1,43 +0,0 @@
-from functools import lru_cache
-from typing import List, Tuple
-
-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 detection_visualize
-from yolo_ import _image_preprocess, _data_postprocess
-
-_HEAD_MODELS = [
-    'head_detect_best_s.onnx',
-    'head_detect_best_n.onnx',
-]
-_DEFAULT_HEAD_MODEL = _HEAD_MODELS[0]
-
-
-@lru_cache()
-def _open_head_detect_model(model_name):
-    return _open_onnx_model(hf_hub_download(
-        'deepghs/imgutils-models',
-        f'head_detect/{model_name}'
-    ))
-
-
-_LABELS = ['head']
-
-
-def detect_heads(image: ImageTyping, model_name: str, max_infer_size=640,
-                 conf_threshold: float = 0.45, iou_threshold: float = 0.7) \
-        -> List[Tuple[Tuple[int, int, int, int], str, float]]:
-    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_head_detect_model(model_name).run(['output0'], {'images': data})
-    return _data_postprocess(output[0], conf_threshold, iou_threshold, old_size, new_size, _LABELS)
-
-
-def _gr_detect_heads(image: ImageTyping, model_name: str, max_infer_size=640,
-                     conf_threshold: float = 0.45, iou_threshold: float = 0.7):
-    ret = detect_heads(image, model_name, max_infer_size, conf_threshold, iou_threshold)
-    return detection_visualize(image, ret, _LABELS)

manbits.py

@@ -1,45 +0,0 @@
-from functools import lru_cache
-from typing import List, Tuple
-
-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 detection_visualize
-from yolo_ import _image_preprocess, _data_postprocess
-
-_MANBIT_MODELS = [
-    'manbits_detect_best_m.onnx',
-]
-_DEFAULT_MANBIT_MODEL = _MANBIT_MODELS[0]
-
-
-@lru_cache()
-def _open_manbits_detect_model(model_name):
-    return _open_onnx_model(hf_hub_download(
-        'deepghs/imgutils-models',
-        f'manbits_detect/{model_name}'
-    ))
-
-
-_LABELS = [
-    'EXPOSED_BELLY', 'EXPOSED_BREAST_F', 'EXPOSED_BREAST_M',
-    'EXPOSED_BUTTOCKS', 'EXPOSED_GENITALIA_F', 'EXPOSED_GENITALIA_M'
-]
-
-
-def detect_manbits(image: ImageTyping, model_name: str, max_infer_size=640,
-                   conf_threshold: float = 0.25, iou_threshold: float = 0.7) \
-        -> List[Tuple[Tuple[int, int, int, int], str, float]]:
-    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_manbits_detect_model(model_name).run(['output0'], {'images': data})
-    return _data_postprocess(output[0], conf_threshold, iou_threshold, old_size, new_size, _LABELS)
-
-
-def _gr_detect_manbits(image: ImageTyping, model_name: str, max_infer_size=640,
-                       conf_threshold: float = 0.25, iou_threshold: float = 0.7):
-    ret = detect_manbits(image, model_name, max_infer_size, conf_threshold, iou_threshold)
-    return detection_visualize(image, ret, _LABELS)

onnx_.py

@@ -1,59 +0,0 @@
-import logging
-import os
-import shutil
-from functools import lru_cache
-from typing import Optional
-
-from hbutils.system import pip_install
-
-
-def _ensure_onnxruntime():
-    try:
-        import onnxruntime
-    except (ImportError, ModuleNotFoundError):
-        logging.warning('Onnx runtime not installed, preparing to install ...')
-        if shutil.which('nvidia-smi'):
-            logging.info('Installing onnxruntime-gpu ...')
-            pip_install(['onnxruntime-gpu'], silent=True)
-        else:
-            logging.info('Installing onnxruntime (cpu) ...')
-            pip_install(['onnxruntime'], silent=True)
-
-
-_ensure_onnxruntime()
-from onnxruntime import get_available_providers, get_all_providers, InferenceSession, SessionOptions, \
-    GraphOptimizationLevel
-
-alias = {
-    'gpu': "CUDAExecutionProvider",
-    "trt": "TensorrtExecutionProvider",
-}
-
-
-def get_onnx_provider(provider: Optional[str] = None):
-    if not provider:
-        if "CUDAExecutionProvider" in get_available_providers():
-            return "CUDAExecutionProvider"
-        else:
-            return "CPUExecutionProvider"
-    elif provider.lower() in alias:
-        return alias[provider.lower()]
-    else:
-        for p in get_all_providers():
-            if provider.lower() == p.lower() or f'{provider}ExecutionProvider'.lower() == p.lower():
-                return p
-
-        raise ValueError(f'One of the {get_all_providers()!r} expected, '
-                         f'but unsupported provider {provider!r} found.')
-
-
-@lru_cache()
-def _open_onnx_model(ckpt: str, provider: str = None) -> InferenceSession:
-    options = SessionOptions()
-    options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL
-    provider = provider or get_onnx_provider()
-    if provider == "CPUExecutionProvider":
-        options.intra_op_num_threads = os.cpu_count()
-
-    logging.info(f'Model {ckpt!r} loaded with provider {provider!r}')
-    return InferenceSession(ckpt, options, [provider])

person.py

@@ -1,46 +0,0 @@
-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 detection_visualize
-from yolo_ import _image_preprocess, _data_postprocess
-
-_PERSON_MODELS = [
-    'person_detect_plus_v1.1_best_m.onnx',
-    'person_detect_plus_v1.1_best_s.onnx',
-    'person_detect_plus_v1.1_best_n.onnx',
-    'person_detect_plus_best_m.onnx',
-    'person_detect_best_m.onnx',
-    'person_detect_best_x.onnx',
-    '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}'
-    ))
-
-
-_LABELS = ['person']
-
-
-def detect_person(image: ImageTyping, model_name: str, max_infer_size=640,
-                  conf_threshold: float = 0.3, iou_threshold: float = 0.5):
-    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_postprocess(output[0], conf_threshold, iou_threshold, old_size, new_size, _LABELS)
-
-
-def _gr_detect_person(image: ImageTyping, model_name: str, max_infer_size=640,
-                      conf_threshold: float = 0.3, iou_threshold: float = 0.5):
-    ret = detect_person(image, model_name, max_infer_size, conf_threshold, iou_threshold)
-    return detection_visualize(image, ret, _LABELS)

plot.py

@@ -1,76 +0,0 @@
-"""
-Overview:
-    Visualize the detection results.
-
-    See :func:`imgutils.detect.face.detect_faces` and :func:`imgutils.detect.person.detect_person` for examples.
-"""
-from typing import List, Tuple, Optional
-
-from PIL import ImageFont, ImageDraw
-from hbutils.color import rnd_colors, Color
-from imgutils.data import ImageTyping, load_image
-
-
-def _try_get_font_from_matplotlib(fontsize: int = 12):
-    try:
-        # noinspection PyPackageRequirements
-        import matplotlib
-    except (ModuleNotFoundError, ImportError):
-        return None
-    else:
-        # noinspection PyPackageRequirements
-        from matplotlib.font_manager import findfont, FontProperties
-        font = findfont(FontProperties(family=['sans-serif']))
-        return ImageFont.truetype(font, fontsize)
-
-
-def detection_visualize(image: ImageTyping, detection: List[Tuple[Tuple[float, float, float, float], str, float]],
-                        labels: Optional[List[str]] = None, text_padding: int = 6, fontsize: int = 12,
-                        no_label: bool = False):
-    """
-    Overview:
-        Visualize the results of the object detection.
-
-    :param image: Image be detected.
-    :param detection: The detection results list, each item includes the detected area `(x0, y0, x1, y1)`,
-        the target type (always `head`) and the target confidence score.
-    :param labels: An array of known labels. If not provided, the labels will be automatically detected
-        from the given ``detection``.
-    :param text_padding: Text padding of the labels. Default is ``6``.
-    :param fontsize: Font size of the labels. At runtime, an attempt will be made to retrieve the font used
-        for rendering from `matplotlib`. Therefore, if `matplotlib` is not installed, only the default pixel font
-        provided with `Pillow` can be used, and the font size cannot be changed.
-    :param no_label: Do not show labels. Default is ``False``.
-    :return: A `PIL` image with the same size as the provided image `image`, which contains the original image
-        content as well as the visualized bounding boxes.
-
-    Examples::
-        See :func:`imgutils.detect.face.detect_faces` and :func:`imgutils.detect.person.detect_person` for examples.
-    """
-    image = load_image(image, force_background=None, mode='RGBA')
-    visual_image = image.copy()
-    draw = ImageDraw.Draw(visual_image, mode='RGBA')
-    font = _try_get_font_from_matplotlib(fontsize) or ImageFont.load_default()
-
-    labels = sorted(labels or {label for _, label, _ in detection})
-    _colors = list(map(str, rnd_colors(len(labels))))
-    _color_map = dict(zip(labels, _colors))
-    for (xmin, ymin, xmax, ymax), label, score in detection:
-        box_color = _color_map[label]
-        draw.rectangle((xmin, ymin, xmax, ymax), outline=box_color, width=2)
-
-        if not no_label:
-            label_text = f'{label}: {score * 100:.2f}%'
-            _t_x0, _t_y0, _t_x1, _t_y1 = draw.textbbox((xmin, ymin), label_text, font=font)
-            _t_width, _t_height = _t_x1 - _t_x0, _t_y1 - _t_y0
-            if ymin - _t_height - text_padding < 0:
-                _t_text_rect = (xmin, ymin, xmin + _t_width + text_padding * 2, ymin + _t_height + text_padding * 2)
-                _t_text_co = (xmin + text_padding, ymin + text_padding)
-            else:
-                _t_text_rect = (xmin, ymin - _t_height - text_padding * 2, xmin + _t_width + text_padding * 2, ymin)
-                _t_text_co = (xmin + text_padding, ymin - _t_height - text_padding)
-
-            draw.rectangle(_t_text_rect, fill=str(Color(box_color, alpha=0.5)))
-            draw.text(_t_text_co, label_text, fill="black", font=font)
-
-    return visual_image

yolo_.py

@@ -1,110 +0,0 @@
-import math
-from typing import List
-
-import numpy as np
-from PIL import Image
-
-
-def _yolo_xywh2xyxy(x: np.ndarray) -> np.ndarray:
-    """
-    Copied from yolov8.
-
-    Convert bounding box coordinates from (x, y, width, height) format to (x1, y1, x2, y2) format where (x1, y1) is the
-    top-left corner and (x2, y2) is the bottom-right corner.
-
-    Args:
-        x (np.ndarray) or (torch.Tensor): The input bounding box coordinates in (x, y, width, height) format.
-    Returns:
-        y (np.ndarray) or (torch.Tensor): The bounding box coordinates in (x1, y1, x2, y2) format.
-    """
-    y = np.copy(x)
-    y[..., 0] = x[..., 0] - x[..., 2] / 2  # top left x
-    y[..., 1] = x[..., 1] - x[..., 3] / 2  # top left y
-    y[..., 2] = x[..., 0] + x[..., 2] / 2  # bottom right x
-    y[..., 3] = x[..., 1] + x[..., 3] / 2  # bottom right y
-    return y
-
-
-def _yolo_nms(boxes, scores, thresh: float = 0.7) -> List[int]:
-    """
-    dets: ndarray, (num_boxes, 5)
-        每一行表示一个bounding box：[xmin, ymin, xmax, ymax, score]
-        其中xmin, ymin, xmax, ymax分别表示框的左上角和右下角坐标，score表示框的分数
-    thresh: float
-        两个框的IoU阈值
-    """
-    x1 = boxes[:, 0]
-    y1 = boxes[:, 1]
-    x2 = boxes[:, 2]
-    y2 = boxes[:, 3]
-    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
-
-    # 按照score降序排列
-    order = scores.argsort()[::-1]
-
-    keep = []
-    while order.size > 0:
-        i = order[0]
-        keep.append(i)
-        # 计算其他所有框与当前框的IoU
-        xx1 = np.maximum(x1[i], x1[order[1:]])
-        yy1 = np.maximum(y1[i], y1[order[1:]])
-        xx2 = np.minimum(x2[i], x2[order[1:]])
-        yy2 = np.minimum(y2[i], y2[order[1:]])
-
-        w = np.maximum(0.0, xx2 - xx1 + 1)
-        h = np.maximum(0.0, yy2 - yy1 + 1)
-
-        inter = w * h
-        iou = inter / (areas[i] + areas[order[1:]] - inter)
-
-        # 保留IoU小于阈值的框
-        inds = np.where(iou <= thresh)[0]
-        order = order[inds + 1]
-
-    return keep
-
-
-def _image_preprocess(image: Image.Image, max_infer_size: int = 640, align: int = 32):
-    old_width, old_height = image.width, image.height
-    new_width, new_height = old_width, old_height
-    r = max_infer_size / max(new_width, new_height)
-    if r < 1:
-        new_width, new_height = new_width * r, new_height * r
-    new_width = int(math.ceil(new_width / align) * align)
-    new_height = int(math.ceil(new_height / align) * align)
-    image = image.resize((new_width, new_height))
-    return image, (old_width, old_height), (new_width, new_height)
-
-
-def _xy_postprocess(x, y, old_size, new_size):
-    old_width, old_height = old_size
-    new_width, new_height = new_size
-    x, y = x / new_width * old_width, y / new_height * old_height
-    x = int(np.clip(x, a_min=0, a_max=old_width).round())
-    y = int(np.clip(y, a_min=0, a_max=old_height).round())
-    return x, y
-
-
-def _data_postprocess(output, conf_threshold, iou_threshold, old_size, new_size, labels: List[str]):
-    max_scores = output[4:, :].max(axis=0)
-    output = output[:, max_scores > conf_threshold].transpose(1, 0)
-    boxes = output[:, :4]
-    scores = output[:, 4:]
-    filtered_max_scores = scores.max(axis=1)
-
-    if not boxes.size:
-        return []
-
-    boxes = _yolo_xywh2xyxy(boxes)
-    idx = _yolo_nms(boxes, filtered_max_scores, thresh=iou_threshold)
-    boxes, scores = boxes[idx], scores[idx]
-
-    detections = []
-    for box, score in zip(boxes, scores):
-        x0, y0 = _xy_postprocess(box[0], box[1], old_size, new_size)
-        x1, y1 = _xy_postprocess(box[2], box[3], old_size, new_size)
-        max_score_id = score.argmax()
-        detections.append(((x0, y0, x1, y1), labels[max_score_id], float(score[max_score_id])))
-
-    return detections