Update app.py

app.py

@@ -32,16 +32,19 @@ 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)
     artist_tags_with_probs = []
     character_tags_with_probs = []
     general_tags_with_probs = []
+    all_artist_tags_probs = [] # Store all artist tags and their probabilities
 
     # Collect tags above thresholds
     for idx, prob in enumerate(probs):
         tag = idx_to_tag[str(idx)]
         cat = tag_to_category.get(tag, "unknown")
+        if cat == 'artist':
+            all_artist_tags_probs.append((tag, float(prob))) # Store all artist tags
         thresh = category_thresholds.get(cat, default_threshold)
         if float(prob) >= thresh:
             # add to category dictionary
@@ -64,6 +67,12 @@ def tag_image(pil_image: Image.Image, output_format: str) -> str:
 
         prompt_tags = artist_prompt_tags + character_prompt_tags + general_prompt_tags
 
+        # Ensure at least one artist tag if any artist tags were predicted at all, even below threshold
+        if not artist_prompt_tags and all_artist_tags_probs:
+            best_artist_tag, best_artist_prob = max(all_artist_tags_probs, key=lambda item: item[1])
+            prompt_tags = [best_artist_tag.replace("_", " ").replace("(", "\\(").replace(")", "\\)")] + prompt_tags
+
+
         if not prompt_tags:
             return "No tags predicted."
         return ", ".join(prompt_tags)