Update

app.py

@@ -31,76 +31,6 @@ def preprocess_image(image):
     image = image.astype(np.float32)
     return image
 
-def process_image(prob, rating_tags, character_tags, general_tags):
-    thresh = 0.35
-
-    # NSFW/SFW判定
-    tag_confidences = {tag: prob[i] for i, tag in enumerate(rating_tags)}
-    max_nsfw_score = max(tag_confidences.get("questionable", 0), tag_confidences.get("explicit", 0))
-    max_sfw_score = tag_confidences.get("general", 0)
-    NSFW_flag = None
-
-    if max_nsfw_score > max_sfw_score:
-        NSFW_flag = "NSFWの可能性が高いです"
-    else:
-        NSFW_flag = "SFWの可能性が高いです"
-
-    # 版権キャラクターの可能性を評価
-    character_tags_with_probs = []
-    for i, p in enumerate(prob[4:]):
-        if p >= thresh and i >= len(general_tags):
-            tag_index = i - len(general_tags)
-            if tag_index < len(character_tags):
-                tag_name = character_tags[tag_index]
-                prob_percent = round(p * 100, 2)  # 確率をパーセンテージに変換
-                character_tags_with_probs.append((tag_name, f"{prob_percent}%"))
-
-    IP_flag = None
-    if character_tags_with_probs:
-        IP_flag = f"版権キャラクター: {character_tags_with_probs}の可能性があります"
-    else:
-        IP_flag = "版権キャラクターの可能性が低いと思われます"
-
-     # タグを生成
-    tag_freq = {}
-    undesired_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)
-
-    # 先頭のカンマを取る
-    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 NSFW_flag, IP_flag, tag_text
-
 
 class webui:
     def __init__(self):
@@ -132,10 +62,75 @@ class webui:
         img = np.array([image])
         prob = ort_sess.run(None, {ort_sess.get_inputs()[0].name: img})[0][0]  # ONNXモデルからの出力
 
-        NSFW_flag, IP_flag, tag_text = process_image(prob, rating_tags, character_tags, general_tags)
+        thresh = 0.35
+
+        # NSFW/SFW判定
+        tag_confidences = {tag: prob[i] for i, tag in enumerate(rating_tags)}
+        max_nsfw_score = max(tag_confidences.get("questionable", 0), tag_confidences.get("explicit", 0))
+        max_sfw_score = tag_confidences.get("general", 0)
+        NSFW_flag = None
+
+        if max_nsfw_score > max_sfw_score:
+            NSFW_flag = "NSFWの可能性が高いです"
+        else:
+            NSFW_flag = "SFWの可能性が高いです"
+
+        # 版権キャラクターの可能性を評価
+        character_tags_with_probs = []
+        for i, p in enumerate(prob[4:]):
+            if p >= thresh and i >= len(general_tags):
+                tag_index = i - len(general_tags)
+                if tag_index < len(character_tags):
+                    tag_name = character_tags[tag_index]
+                    prob_percent = round(p * 100, 2)  # 確率をパーセンテージに変換
+                    character_tags_with_probs.append((tag_name, f"{prob_percent}%"))
+
+        IP_flag = None
+        if character_tags_with_probs:
+            IP_flag = f"版権キャラクター: {character_tags_with_probs}の可能性があります"
+        else:
+            IP_flag = "版権キャラクターの可能性が低いと思われます"
+
+        # タグを生成
+        tag_freq = {}
+        undesired_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)
+
+        # 先頭のカンマを取る
+        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 NSFW_flag, IP_flag, tag_text
 
-
     def launch(self):
         with self.demo:
             with gr.Row():