dev(narugo): add new models

aicheck.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_ai_check'
+_AICHECK_MODELS = natsorted([
+    os.path.dirname(os.path.relpath(file, _REPO))
+    for file in hfs.glob(f'{_REPO}/*/model.onnx')
+])
+_DEFAULT_AICHECK_MODEL = 'mobilenetv3_sce_dist'
+
+
+@lru_cache()
+def _open_anime_aicheck_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_aicheck(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_aicheck_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

app.py

@@ -2,8 +2,10 @@ import os
 
 import gradio as gr
 
+from aicheck import _gr_aicheck, _DEFAULT_AICHECK_MODEL, _AICHECK_MODELS
 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
 
 if __name__ == '__main__':
     with gr.Blocks() as demo:
@@ -42,4 +44,38 @@ if __name__ == '__main__':
                         outputs=[gr_mono_output],
                     )
 
+            with gr.Tab('AI Check'):
+                with gr.Row():
+                    with gr.Column():
+                        gr_aicheck_input_image = gr.Image(type='pil', label='Original Image')
+                        gr_aicheck_model = gr.Dropdown(_AICHECK_MODELS, value=_DEFAULT_AICHECK_MODEL, label='Model')
+                        gr_aicheck_infer_size = gr.Slider(224, 640, value=384, step=32, label='Infer Size')
+                        gr_aicheck_submit = gr.Button(value='Submit', variant='primary')
+
+                    with gr.Column():
+                        gr_aicheck_output = gr.Label(label='Classes')
+
+                    gr_aicheck_submit.click(
+                        _gr_aicheck,
+                        inputs=[gr_aicheck_input_image, gr_aicheck_model, gr_aicheck_infer_size],
+                        outputs=[gr_aicheck_output],
+                    )
+
+            with gr.Tab('Rating'):
+                with gr.Row():
+                    with gr.Column():
+                        gr_rating_input_image = gr.Image(type='pil', label='Original Image')
+                        gr_rating_model = gr.Dropdown(_RATING_MODELS, value=_DEFAULT_RATING_MODEL, label='Model')
+                        gr_rating_infer_size = gr.Slider(224, 640, value=384, step=32, label='Infer Size')
+                        gr_rating_submit = gr.Button(value='Submit', variant='primary')
+
+                    with gr.Column():
+                        gr_rating_output = gr.Label(label='Classes')
+
+                    gr_rating_submit.click(
+                        _gr_rating,
+                        inputs=[gr_rating_input_image, gr_rating_model, gr_rating_infer_size],
+                        outputs=[gr_rating_output],
+                    )
+
     demo.queue(os.cpu_count()).launch()

rating.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_rating'
+_RATING_MODELS = natsorted([
+    os.path.dirname(os.path.relpath(file, _REPO))
+    for file in hfs.glob(f'{_REPO}/*/model.onnx')
+])
+_DEFAULT_RATING_MODEL = 'mobilenetv3_sce_dist'
+
+
+@lru_cache()
+def _open_anime_rating_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_rating(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_rating_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