deepghs

/

anime_real_cls

like 6

Follow

DeepGHS 336

inference code:

import json
import numpy as np
from PIL import Image
from imgutils.data import load_image, rgb_encode
from onnxruntime import InferenceSession, SessionOptions, GraphOptimizationLevel

class Anime_Real_Cls():
    def __init__(self, model_dir):
        model_path = f'{model_dir}/model.onnx'
        self.model = self.load_local_onnx_model(model_path)
        with open(f'{model_dir}/meta.json', 'r') as f:
            self.labels = json.load(f)['labels']

    def _img_encode(self, image_path, size=(384, 384), normalize=(0.5, 0.5)):
        image = Image.open(image_path)
        image = load_image(image, mode='RGB')
        image = image.resize(size, Image.BILINEAR)
        data = rgb_encode(image, order_='CHW')
        if normalize:
            mean_, std_ = normalize
            mean = np.asarray([mean_]).reshape((-1, 1, 1))
            std = np.asarray([std_]).reshape((-1, 1, 1))
            data = (data - mean) / std
        return data.astype(np.float32)

    def load_local_onnx_model(self, model_path: str) -> InferenceSession:
        options = SessionOptions()
        options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL
        return InferenceSession(model_path, options)

    def __call__(self, image_path):
        input_ = self._img_encode(image_path, size=(384, 384))[None, ...]
        output, = self.model.run(['output'], {'input': input_})
        values = dict(zip(self.labels, map(lambda x: x.item(), output[0])))
        print("values: ", values)
        max_key = max(values, key=values.get)
        return max_key

if __name__ == "__main__":
    classifier = Anime_Real_Cls(model_dir="./caformer_s36_v1.3_fixed")
    image_path = '1.webp'
    class_result = classifier(image_path)
    print("class_result: ", class_result)