Update README.md

README.md

@@ -7,5 +7,39 @@ Model convert from [https://github.com/KichangKim/DeepDanbooru](https://github.c
 Usage:
 
 ```python
+import cv2
+import numpy as np
+import onnxruntime as rt
+from huggingface_hub import hf_hub_download
 
+tagger_model_path = hf_hub_download(repo_id="skytnt/deepdanbooru_onnx", filename="deepdanbooru.onnx")
+
+tagger_model = rt.InferenceSession(tagger_model_path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
+tagger_model_meta = tagger_model.get_modelmeta().custom_metadata_map
+tagger_tags = eval(tagger_model_meta['tags'])
+
+def tagger_predict(image, score_threshold):
+        h, w = image.shape[:2]
+        r = min(512 / w, 512 / h)
+        h, w = int(h * r), int(w * r)
+        image = cv2.resize(image, (w, h))
+        pdx = 512 - w
+        pdy = 512 - h
+        img_new = np.full([512, 512, 3], 1, dtype=np.float32)
+        img_new[pdy // 2:pdy // 2 + h, pdx // 2:pdx // 2 + w] = image
+        image = img_new[np.newaxis, :]
+        probs = tagger_model.run(None, {"input_1": image})[0][0]
+        probs = probs.astype(np.float32)
+        res = []
+        for prob, label in zip(probs.tolist(), tagger_tags):
+            if prob < score_threshold:
+                continue
+            res.append(label)
+        return res
+
+img = cv2.imread("test.jpg")
+img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+img = img.astype(np.float32) / 255
+tags = tagger_predict(img, 0.5)
+print(tags)
 ```