dev(narugo): add censor point detection

app.py

@@ -2,6 +2,7 @@ import os
 
 import gradio as gr
 
+from censor import _CENSOR_MODELS, _DEFAULT_CENSOR_MODEL, _gr_detect_censors
 from face import _FACE_MODELS, _DEFAULT_FACE_MODEL, _gr_detect_faces
 from head import _gr_detect_heads, _HEAD_MODELS, _DEFAULT_HEAD_MODEL
 from manbits import _MANBIT_MODELS, _DEFAULT_MANBIT_MODEL, _gr_detect_manbits
@@ -18,7 +19,7 @@ if __name__ == '__main__':
                         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.45, label='Score 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')
 
@@ -82,6 +83,30 @@ if __name__ == '__main__':
                         outputs=[gr_person_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.35, 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'):
                 with gr.Row():
                     with gr.Column():

censor.py

@@ -0,0 +1,42 @@
+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_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.35, 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.35, iou_threshold: float = 0.5):
+    ret = detect_censors(image, model_name, max_infer_size, conf_threshold, iou_threshold)
+    return detection_visualize(image, ret, _LABELS)

face.py

@@ -26,7 +26,7 @@ _DEFAULT_FACE_MODEL = _FACE_MODELS[0]
 def _open_face_detect_model(model_name):
     return _open_onnx_model(hf_hub_download(
         f'deepghs/anime_face_detection',
-        f'{model_name}/model.onnx'
+        f'{model_name}/model.onnx',
     ))
 
 
@@ -34,7 +34,7 @@ _LABELS = ['face']
 
 
 def detect_faces(image: ImageTyping, model_name: str, max_infer_size=640,
-                 conf_threshold: float = 0.45, iou_threshold: float = 0.7) \
+                 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)
@@ -45,6 +45,6 @@ def detect_faces(image: ImageTyping, model_name: str, max_infer_size=640,
 
 
 def _gr_detect_faces(image: ImageTyping, model_name: str, max_infer_size=640,
-                     conf_threshold: float = 0.45, iou_threshold: float = 0.7):
+                     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)