Update app.py

app.py

@@ -44,13 +44,15 @@ def run_inference(pil_image: Image.Image) -> np.ndarray:
     _, refined_logits = session.run(None, {input_name: input_tensor})
     return refined_logits[0]
 
-def get_tags(refined_logits: np.ndarray, metadata: dict, default_threshold: float = DEFAULT_THRESHOLD):
+def get_tags(refined_logits: np.ndarray, metadata: dict, custom_threshold: float = None):
     """
     Compute probabilities from logits and collect tag predictions.
     
+    If custom_threshold is provided, it overrides category-specific thresholds.
+    
     Returns:
         results_by_cat: Dictionary mapping each category to a list of (tag, probability) above its threshold.
-        prompt_tags_by_cat: Similar dictionary but only for prompt-style categories (artist, character, general).
+        prompt_tags_by_cat: Dictionary for prompt-style output with keys: artist, character, general.
         all_artist_tags: All artist tags (with probabilities) regardless of threshold.
     """
     probs = 1 / (1 + np.exp(-refined_logits))
@@ -65,7 +67,8 @@ def get_tags(refined_logits: np.ndarray, metadata: dict, default_threshold: floa
     for idx, prob in enumerate(probs):
         tag = idx_to_tag[str(idx)]
         cat = tag_to_category.get(tag, "unknown")
-        thresh = category_thresholds.get(cat, default_threshold)
+        # Use custom threshold if provided; otherwise, use metadata threshold or default.
+        thresh = custom_threshold if custom_threshold is not None else category_thresholds.get(cat, DEFAULT_THRESHOLD)
         if cat == "artist":
             all_artist_tags.append((tag, float(prob)))
         if float(prob) >= thresh:
@@ -80,7 +83,6 @@ def format_prompt_tags(prompt_tags_by_cat: dict, all_artist_tags: list) -> str:
     
     Returns a comma-separated string of escaped tags.
     """
-    # Sort tags within each category by probability (descending)
     for cat in prompt_tags_by_cat:
         prompt_tags_by_cat[cat].sort(key=lambda x: x[1], reverse=True)
     
@@ -89,7 +91,7 @@ def format_prompt_tags(prompt_tags_by_cat: dict, all_artist_tags: list) -> str:
     general_tags = [escape_tag(tag) for tag, _ in prompt_tags_by_cat.get("general", [])]
     prompt_tags = artist_tags + character_tags + general_tags
 
-    # Ensure at least one artist tag appears if available, even if below threshold
+    # Ensure at least one artist tag appears even if none pass the threshold
     if not artist_tags and all_artist_tags:
         best_artist_tag, _ = max(all_artist_tags, key=lambda item: item[1])
         prompt_tags.insert(0, escape_tag(best_artist_tag))
@@ -115,16 +117,20 @@ def format_detailed_output(results_by_cat: dict, all_artist_tags: list) -> str:
         lines.append(f"**Category: {cat}** – {len(tag_list)} tags")
         for tag, prob in tag_list:
             lines.append(f"- {escape_tag(tag)} (Prob: {prob:.3f})")
-        lines.append("")  # blank line between categories
+        lines.append("") 
     return "\n".join(lines)
 
-def tag_image(pil_image: Image.Image, output_format: str) -> str:
-    """Run inference on the image and return formatted tags based on the chosen output format."""
+def tag_image(pil_image: Image.Image, output_format: str, threshold: float) -> str:
+    """
+    Run inference on the image and return formatted tags based on the chosen output format.
+    
+    The threshold slider value overrides category-specific thresholds if provided.
+    """
     if pil_image is None:
         return "Please upload an image."
     
     refined_logits = run_inference(pil_image)
-    results_by_cat, prompt_tags_by_cat, all_artist_tags = get_tags(refined_logits, metadata)
+    results_by_cat, prompt_tags_by_cat, all_artist_tags = get_tags(refined_logits, metadata, custom_threshold=threshold)
     
     if output_format == "Prompt-style Tags":
         return format_prompt_tags(prompt_tags_by_cat, all_artist_tags)
@@ -152,11 +158,23 @@ with demo:
                 value="Prompt-style Tags",
                 label="Output Format"
             )
+            # Slider to modify the global threshold value
+            threshold_slider = gr.Slider(
+                minimum=0,
+                maximum=1,
+                step=0.05,
+                value=DEFAULT_THRESHOLD,
+                label="Global Threshold"
+            )
             tag_button = gr.Button("🔍 Tag Image")
         with gr.Column():
-            output_box = gr.Markdown("")  # Markdown output for formatted results
+            output_box = gr.Markdown("")
 
-    tag_button.click(fn=tag_image, inputs=[image_in, format_choice], outputs=output_box)
+    tag_button.click(
+        fn=tag_image,
+        inputs=[image_in, format_choice, threshold_slider],
+        outputs=output_box
+    )
     
     gr.Markdown(
         "----\n"