dev(narugo): add monochrome

app.py

@@ -3,6 +3,7 @@ import os
 import gradio as gr
 
 from cls import _CLS_MODELS, _DEFAULT_CLS_MODEL, _gr_classification
+from monochrome import _gr_monochrome, _DEFAULT_MONO_MODEL, _MONO_MODELS
 
 if __name__ == '__main__':
     with gr.Blocks() as demo:
@@ -24,4 +25,21 @@ if __name__ == '__main__':
                         outputs=[gr_cls_output],
                     )
 
+            with gr.Tab('Monochrome'):
+                with gr.Row():
+                    with gr.Column():
+                        gr_mono_input_image = gr.Image(type='pil', label='Original Image')
+                        gr_mono_model = gr.Dropdown(_MONO_MODELS, value=_DEFAULT_MONO_MODEL, label='Model')
+                        gr_mono_infer_size = gr.Slider(224, 640, value=384, step=32, label='Infer Size')
+                        gr_mono_submit = gr.Button(value='Submit', variant='primary')
+
+                    with gr.Column():
+                        gr_mono_output = gr.Label(label='Classes')
+
+                    gr_mono_submit.click(
+                        _gr_monochrome,
+                        inputs=[gr_mono_input_image, gr_mono_model, gr_mono_infer_size],
+                        outputs=[gr_mono_output],
+                    )
+
     demo.queue(os.cpu_count()).launch()

cls.py

@@ -1,31 +1,34 @@
+import json
+import os
 from functools import lru_cache
-from typing import Mapping
+from typing import Mapping, List
 
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, HfFileSystem
 from imgutils.data import ImageTyping, load_image
+from natsort import natsorted
 
 from onnx_ import _open_onnx_model
 from preprocess import _img_encode
 
-_LABELS = ['3d', 'bangumi', 'comic', 'illustration']
-_CLS_MODELS = [
-    'caformer_s36',
-    'caformer_s36_plus',
-    'mobilenetv3',
-    'mobilenetv3_dist',
-    'mobilenetv3_sce',
-    'mobilenetv3_sce_dist',
-    'mobilevitv2_150',
-]
+hfs = HfFileSystem()
+
+_REPO = 'deepghs/anime_classification'
+_CLS_MODELS = natsorted([
+    os.path.dirname(os.path.relpath(file, _REPO))
+    for file in hfs.glob(f'{_REPO}/*/model.onnx')
+])
 _DEFAULT_CLS_MODEL = 'mobilenetv3_sce_dist'
 
 
 @lru_cache()
 def _open_anime_classify_model(model_name):
-    return _open_onnx_model(hf_hub_download(
-        f'deepghs/anime_classification',
-        f'{model_name}/model.onnx',
-    ))
+    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_classification(image: ImageTyping, model_name: str, size=384) -> Mapping[str, float]:
@@ -33,5 +36,6 @@ def _gr_classification(image: ImageTyping, model_name: str, size=384) -> Mapping
     input_ = _img_encode(image, size=(size, size))[None, ...]
     output, = _open_anime_classify_model(model_name).run(['output'], {'input': input_})
 
-    values = dict(zip(_LABELS, map(lambda x: x.item(), output[0])))
+    labels = _get_tags(model_name)
+    values = dict(zip(labels, map(lambda x: x.item(), output[0])))
     return values

monochrome.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/monochrome_detect'
+_MONO_MODELS = natsorted([
+    os.path.dirname(os.path.relpath(file, _REPO))
+    for file in hfs.glob(f'{_REPO}/*/model.onnx')
+])
+_DEFAULT_MONO_MODEL = 'mobilenetv3_large_100_dist'
+
+
+@lru_cache()
+def _open_anime_monochrome_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_monochrome(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_monochrome_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

requirements.txt

@@ -2,10 +2,11 @@ gradio==3.18.0
 numpy
 pillow
 onnxruntime
-huggingface_hub
+huggingface_hub>=0.14.0
 scikit-image
 pandas
 opencv-python>=4.6.0
 hbutils>=0.9.0
 dghs-imgutils>=0.1.0
 httpx==0.23.0
+natsort