Update app.py

app.py

@@ -32,12 +32,13 @@ def tag_image(pil_image: Image.Image, output_format: str) -> str:
     idx_to_tag = metadata["idx_to_tag"]
     tag_to_category = metadata.get("tag_to_category", {})
     category_thresholds = metadata.get("category_thresholds", {})
-    default_threshold = 0.325
+    default_threshold = 0.35
     results_by_cat = {}  # to store tags per category (for verbose output)
-    prompt_tags = []     # to store tags for prompt-style output
-    # Collect tags above thresholds
     artist_tags_with_probs = []
-    non_artist_tags_with_probs = []
+    character_tags_with_probs = []
+    general_tags_with_probs = []
+
+    # Collect tags above thresholds
     for idx, prob in enumerate(probs):
         tag = idx_to_tag[str(idx)]
         cat = tag_to_category.get(tag, "unknown")
@@ -47,16 +48,21 @@ def tag_image(pil_image: Image.Image, output_format: str) -> str:
             results_by_cat.setdefault(cat, []).append((tag, float(prob)))
             if cat == 'artist':
                 artist_tags_with_probs.append((tag, float(prob)))
-            else:
-                non_artist_tags_with_probs.append((tag, float(prob)))
+            elif cat == 'character':
+                character_tags_with_probs.append((tag, float(prob)))
+            elif cat == 'general':
+                general_tags_with_probs.append((tag, float(prob)))
 
     if output_format == "Prompt-style Tags":
         artist_tags_with_probs.sort(key=lambda x: x[1], reverse=True)
-        non_artist_tags_with_probs.sort(key=lambda x: x[1], reverse=True)
+        character_tags_with_probs.sort(key=lambda x: x[1], reverse=True)
+        general_tags_with_probs.sort(key=lambda x: x[1], reverse=True)
+
+        artist_prompt_tags = [tag.replace("_", " ").replace("(", "\\(").replace(")", "\\)") for tag, prob in artist_tags_with_probs]
+        character_prompt_tags = [tag.replace("_", " ").replace("(", "\\(").replace(")", "\\)") for tag, prob in character_tags_with_probs]
+        general_prompt_tags = [tag.replace("_", " ").replace("(", "\\(").replace(")", "\\)") for tag, prob in general_tags_with_probs]
 
-        artist_prompt_tags = [tag.replace("_", " ") for tag, prob in artist_tags_with_probs]
-        non_artist_prompt_tags = [tag.replace("_", " ") for tag, prob in non_artist_tags_with_probs]
-        prompt_tags = artist_prompt_tags + non_artist_prompt_tags
+        prompt_tags = artist_prompt_tags + character_prompt_tags + general_prompt_tags
 
         if not prompt_tags:
             return "No tags predicted."