Upload helper.py

sd/prompt_helper/helper.py

@@ -1,11 +1,19 @@
 import onnx
 import onnxruntime as ort
 import numpy as np
+import cv2
+import os
+import csv
 
 VIT_MODEL_DSV3_REPO = "SmilingWolf/wd-vit-tagger-v3"
 
+MODEL_DIR='sd/prompt_helper/model'
+ORT_SESSION=ort.InferenceSession(f'{MODEL_DIR}/model.onnx', providers=['CPUExecutionProvider'])
+#ORT_INPUT_NAME=ort_session.get_inputs()[0].name
 
-def preprocess_image(img):
+IMAGE_SIZE = 448
+
+def img_preporation(img):
     bgr_img = np.array(img)[:, :, ::-1].copy()
     size = max(bgr_img.shape[0:2])
     pad_x = size - bgr_img.shape[1]
@@ -20,4 +28,65 @@ def preprocess_image(img):
     interp = cv2.INTER_AREA if size > IMAGE_SIZE else cv2.INTER_LANCZOS4
     bgr_img = cv2.resize(bgr_img, (IMAGE_SIZE, IMAGE_SIZE), interpolation=interp)
 
-    bgr_img = bgr_img.astype(np.float32)
+    bgr_img = bgr_img.astype(np.float32)
+
+
+def analysis(img):
+    # = model[0]
+    ort_input_name = ORT_SESSION.get_inputs()[0].name
+
+    with open(os.path.join(MODEL_DIR, "selected_tags.csv"), "r", encoding="utf-8") as f:
+        reader = csv.reader(f)
+        l = [row for row in reader]
+        header = l[0]  # tag_id,name,category,count
+        rows = l[1:]
+    assert header[0] == "tag_id" and header[1] == "name" and header[2] == "category", f"unexpected csv format: {header}"
+
+    general_tags = [row[1] for row in rows[1:] if row[2] == "0"]
+    character_tags = [row[1] for row in rows[1:] if row[2] == "4"]
+
+    tag_freq = {}
+    undesired_tags = ["transparent background"]
+
+    #img = Image.open(image_path)
+    preped_img = img_preporation(img)
+    preped_img = np.expand_dims(preped_img, axis=0)
+
+    # Run inference
+    prob = ORT_SESSION.run(None, {ort_input_name: preped_img})[0][0]
+    # Generate Tags
+    combined_tags = []
+    general_tag_text = ""
+    character_tag_text = ""
+    remove_underscore = True
+    caption_separator = ", "
+    general_threshold = 0.35
+    character_threshold = 0.35
+
+    for i, p in enumerate(prob[4:]):
+        if i < len(general_tags) and p >= general_threshold:
+            tag_name = general_tags[i]
+            if remove_underscore and len(tag_name) > 3:  # ignore emoji tags like >_< and ^_^
+                tag_name = tag_name.replace("_", " ")
+
+            if tag_name not in undesired_tags:
+                tag_freq[tag_name] = tag_freq.get(tag_name, 0) + 1
+                general_tag_text += caption_separator + tag_name
+                combined_tags.append(tag_name)
+        elif i >= len(general_tags) and p >= character_threshold:
+            tag_name = character_tags[i - len(general_tags)]
+            if remove_underscore and len(tag_name) > 3:
+                tag_name = tag_name.replace("_", " ")
+
+            if tag_name not in undesired_tags:
+                tag_freq[tag_name] = tag_freq.get(tag_name, 0) + 1
+                character_tag_text += caption_separator + tag_name
+                combined_tags.append(tag_name)
+
+    # Remove leading comma
+    if len(general_tag_text) > 0:
+        general_tag_text = general_tag_text[len(caption_separator) :]
+    if len(character_tag_text) > 0:
+        character_tag_text = character_tag_text[len(caption_separator) :]
+    tag_text = caption_separator.join(combined_tags)
+    return tag_text