Update README.md

README.md

@@ -28,3 +28,50 @@ tags:
 |     mobilenetv3_v0_dist     |  0.63G  |  4.18M   |   99.14%   | 0.9986 |     [confusion](https://huggingface.co/deepghs/anime_real_cls/blob/main/mobilenetv3_v0_dist/plot_confusion.png)     | `anime`, `real` |
 |       caformer_s36_v0       | 22.10G  |  37.21M  |   99.34%   | 0.9988 |       [confusion](https://huggingface.co/deepghs/anime_real_cls/blob/main/caformer_s36_v0/plot_confusion.png)       | `anime`, `real` |
 
+
+```
+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)
+
+```