dev(narugo): upload chsex models

app.py

@@ -3,6 +3,7 @@ import os
 import gradio as gr
 
 from aicheck import _gr_aicheck, _DEFAULT_AICHECK_MODEL, _AICHECK_MODELS
+from chsex import _gr_chsex, _CHSEX_MODELS, _DEFAULT_CHSEX_MODEL
 from cls import _CLS_MODELS, _DEFAULT_CLS_MODEL, _gr_classification
 from monochrome import _gr_monochrome, _DEFAULT_MONO_MODEL, _MONO_MODELS
 from rating import _RATING_MODELS, _DEFAULT_RATING_MODEL, _gr_rating
@@ -78,4 +79,21 @@ if __name__ == '__main__':
                         outputs=[gr_rating_output],
                     )
 
+            with gr.Tab('Character Sex'):
+                with gr.Row():
+                    with gr.Column():
+                        gr_chsex_input_image = gr.Image(type='pil', label='Original Image')
+                        gr_chsex_model = gr.Dropdown(_CHSEX_MODELS, value=_DEFAULT_CHSEX_MODEL, label='Model')
+                        gr_chsex_infer_size = gr.Slider(224, 640, value=384, step=32, label='Infer Size')
+                        gr_chsex_submit = gr.Button(value='Submit', variant='primary')
+
+                    with gr.Column():
+                        gr_chsex_output = gr.Label(label='Classes')
+
+                    gr_chsex_submit.click(
+                        _gr_chsex,
+                        inputs=[gr_chsex_input_image, gr_chsex_model, gr_chsex_infer_size],
+                        outputs=[gr_chsex_output],
+                    )
+
     demo.queue(os.cpu_count()).launch()

chsex.py

@@ -0,0 +1,42 @@
+import json
+import os
+from functools import lru_cache
+from typing import Mapping, List
+
+from huggingface_hub import HfFileSystem
+from huggingface_hub import hf_hub_download
+from imgutils.data import ImageTyping, load_image
+from natsort import natsorted
+
+from onnx_ import _open_onnx_model
+from preprocess import _img_encode
+
+hfs = HfFileSystem()
+
+_REPO = 'deepghs/anime_ch_sex'
+_CHSEX_MODELS = natsorted([
+    os.path.dirname(os.path.relpath(file, _REPO))
+    for file in hfs.glob(f'{_REPO}/*/model.onnx')
+])
+_DEFAULT_CHSEX_MODEL = 'caformer_s36_v1'
+
+
+@lru_cache()
+def _open_anime_chsex_model(model_name):
+    return _open_onnx_model(hf_hub_download(_REPO, f'{model_name}/model.onnx'))
+
+
+@lru_cache()
+def _get_tags(model_name) -> List[str]:
+    with open(hf_hub_download(_REPO, f'{model_name}/meta.json'), 'r') as f:
+        return json.load(f)['labels']
+
+
+def _gr_chsex(image: ImageTyping, model_name: str, size=384) -> Mapping[str, float]:
+    image = load_image(image, mode='RGB')
+    input_ = _img_encode(image, size=(size, size))[None, ...]
+    output, = _open_anime_chsex_model(model_name).run(['output'], {'input': input_})
+
+    labels = _get_tags(model_name)
+    values = dict(zip(labels, map(lambda x: x.item(), output[0])))
+    return values