Update app.py

app.py

@@ -10,7 +10,7 @@ MODEL_REPO = "AngelBottomless/camie-tagger-onnxruntime"
 MODEL_FILE = "camie_tagger_initial.onnx"
 META_FILE = "metadata.json"
 IMAGE_SIZE = (512, 512)
-DEFAULT_THRESHOLD = 0.35
+DEFAULT_THRESHOLD = 0.35  # Default value if slider is not used
 
 # Download model and metadata from Hugging Face Hub
 model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE, cache_dir=".")
@@ -44,15 +44,13 @@ 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, custom_threshold: float = None):
+def get_tags(refined_logits: np.ndarray, metadata: dict, default_threshold: float):
     """
     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: Dictionary for prompt-style output with keys: artist, character, general.
+        prompt_tags_by_cat: Dictionary for prompt-style output (artist, character, general).
         all_artist_tags: All artist tags (with probabilities) regardless of threshold.
     """
     probs = 1 / (1 + np.exp(-refined_logits))
@@ -67,8 +65,7 @@ def get_tags(refined_logits: np.ndarray, metadata: dict, custom_threshold: float
     for idx, prob in enumerate(probs):
         tag = idx_to_tag[str(idx)]
         cat = tag_to_category.get(tag, "unknown")
-        # 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)
+        thresh = category_thresholds.get(cat, default_threshold)
         if cat == "artist":
             all_artist_tags.append((tag, float(prob)))
         if float(prob) >= thresh:
@@ -83,6 +80,7 @@ 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)
     
@@ -91,7 +89,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 even if none pass the threshold
+    # Ensure at least one artist tag appears if available, even if below 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))
@@ -117,20 +115,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("") 
+        lines.append("")  # blank line between categories
     return "\n".join(lines)
 
 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.
+    The slider value (threshold) overrides the default threshold for tag selection.
     """
     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, custom_threshold=threshold)
+    results_by_cat, prompt_tags_by_cat, all_artist_tags = get_tags(refined_logits, metadata, default_threshold=threshold)
     
     if output_format == "Prompt-style Tags":
         return format_prompt_tags(prompt_tags_by_cat, all_artist_tags)
@@ -144,10 +142,10 @@ with demo:
     gr.Markdown(
         "# 🏷️ Camie Tagger – Anime Image Tagging\n"
         "This demo uses an ONNX model of Camie Tagger to label anime illustrations with tags. "
-        "Upload an image and click **Tag Image** to see predictions."
+        "Upload an image, adjust the threshold, and click **Tag Image** to see predictions."
     )
     gr.Markdown(
-        "*(Note: The model will predict a large number of tags across categories like character, general, artist, etc. "
+        "*(Note: The model predicts a large number of tags across categories like character, general, artist, etc. "
         "You can choose a concise prompt-style output or a detailed category-wise breakdown.)*"
     )
     with gr.Row():
@@ -158,23 +156,20 @@ with demo:
                 value="Prompt-style Tags",
                 label="Output Format"
             )
-            # Slider to modify the global threshold value
+            # Slider to modify the default threshold value used in inference.
             threshold_slider = gr.Slider(
-                minimum=0,
-                maximum=1,
+                minimum=0.0,
+                maximum=1.0,
                 step=0.05,
                 value=DEFAULT_THRESHOLD,
-                label="Global Threshold"
+                label="Threshold"
             )
             tag_button = gr.Button("🔍 Tag Image")
         with gr.Column():
-            output_box = gr.Markdown("")
+            output_box = gr.Markdown("")  # Markdown output for formatted results
 
-    tag_button.click(
-        fn=tag_image,
-        inputs=[image_in, format_choice, threshold_slider],
-        outputs=output_box
-    )
+    # Pass the threshold_slider value into the tag_image function
+    tag_button.click(fn=tag_image, inputs=[image_in, format_choice, threshold_slider], outputs=output_box)
     
     gr.Markdown(
         "----\n"