dev(narugo): simplify code

aicheck.py

@@ -1,42 +0,0 @@
-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,98 +2,21 @@ 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
+from base import Classification
+
+apps = [
+    Classification('Classification', 'deepghs/anime_classification', 'mobilenetv3_sce_dist'),
+    Classification('Monochrome', 'deepghs/monochrome_detect', 'mobilenetv3_large_100_dist'),
+    Classification('AI Check', 'deepghs/anime_ai_check', 'mobilenetv3_sce_dist'),
+    Classification('Rating', 'deepghs/anime_rating', 'mobilenetv3_sce_dist'),
+    Classification('Character Sex', 'deepghs/anime_ch_sex', 'caformer_s36_v1'),
+    Classification('Character Skin', 'deepghs/anime_ch_skin_color', 'caformer_s36'),
+]
 
 if __name__ == '__main__':
     with gr.Blocks() as demo:
         with gr.Tabs():
-            with gr.Tab('Classification'):
-                with gr.Row():
-                    with gr.Column():
-                        gr_cls_input_image = gr.Image(type='pil', label='Original Image')
-                        gr_cls_model = gr.Dropdown(_CLS_MODELS, value=_DEFAULT_CLS_MODEL, label='Model')
-                        gr_cls_infer_size = gr.Slider(224, 640, value=384, step=32, label='Infer Size')
-                        gr_cls_submit = gr.Button(value='Submit', variant='primary')
-
-                    with gr.Column():
-                        gr_cls_output = gr.Label(label='Classes')
-
-                    gr_cls_submit.click(
-                        _gr_classification,
-                        inputs=[gr_cls_input_image, gr_cls_model, gr_cls_infer_size],
-                        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],
-                    )
-
-            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],
-                    )
-
-            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],
-                    )
+            for cls in apps:
+                cls.create_gr()
 
     demo.queue(os.cpu_count()).launch()

base.py

@@ -0,0 +1,65 @@
+import json
+import os
+from functools import lru_cache
+from typing import Mapping
+
+import gradio as gr
+from huggingface_hub import HfFileSystem, 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()
+
+
+@lru_cache()
+def open_model_from_repo(repository, model):
+    runtime = _open_onnx_model(hf_hub_download(repository, f'{model}/model.onnx'))
+    with open(hf_hub_download(repository, f'{model}/meta.json'), 'r') as f:
+        labels = json.load(f)['labels']
+
+    return runtime, labels
+
+
+class Classification:
+    def __init__(self, title: str, repository: str, default_model=None, imgsize: int = 384):
+        self.title = title
+        self.repository = repository
+        self.models = natsorted([
+            os.path.dirname(os.path.relpath(file, self.repository))
+            for file in hfs.glob(f'{self.repository}/*/model.onnx')
+        ])
+        self.default_model = default_model or self.models[0]
+        self.imgsize = imgsize
+
+    def _open_onnx_model(self, model):
+        return open_model_from_repo(self.repository, model)
+
+    def _gr_classification(self, image: ImageTyping, model_name: str, size=384) -> Mapping[str, float]:
+        image = load_image(image, mode='RGB')
+        input_ = _img_encode(image, size=(size, size))[None, ...]
+        model, labels = self._open_onnx_model(model_name)
+        output, = model.run(['output'], {'input': input_})
+
+        values = dict(zip(labels, map(lambda x: x.item(), output[0])))
+        return values
+
+    def create_gr(self):
+        with gr.Tab(self.title):
+            with gr.Row():
+                with gr.Column():
+                    gr_input_image = gr.Image(type='pil', label='Original Image')
+                    gr_model = gr.Dropdown(self.models, value=self.default_model, label='Model')
+                    gr_infer_size = gr.Slider(224, 640, value=384, step=32, label='Infer Size')
+                    gr_submit = gr.Button(value='Submit', variant='primary')
+
+                with gr.Column():
+                    gr_output = gr.Label(label='Classes')
+
+                gr_submit.click(
+                    self._gr_classification,
+                    inputs=[gr_input_image, gr_model, gr_infer_size],
+                    outputs=[gr_output],
+                )

chsex.py

@@ -1,42 +0,0 @@
-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

cls.py

@@ -1,41 +0,0 @@
-import json
-import os
-from functools import lru_cache
-from typing import Mapping, List
-
-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
-
-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(_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]:
-    image = load_image(image, mode='RGB')
-    input_ = _img_encode(image, size=(size, size))[None, ...]
-    output, = _open_anime_classify_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

monochrome.py

@@ -1,42 +0,0 @@
-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

rating.py

@@ -1,42 +0,0 @@
-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