cls.py

from functools import lru_cache
from typing import Mapping

from huggingface_hub import hf_hub_download
from imgutils.data import ImageTyping, load_image

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',
]
_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',
    ))


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_})

    values = dict(zip(_LABELS, map(lambda x: x.item(), output[0])))
    return values