Update app.py

app.py

@@ -5,120 +5,165 @@ from PIL import Image
 import json
 from huggingface_hub import hf_hub_download
 
-# Load model and metadata at startup (same as before)
+# Constants
 MODEL_REPO = "AngelBottomless/camie-tagger-onnxruntime"
 MODEL_FILE = "camie_tagger_initial.onnx"
 META_FILE = "metadata.json"
+IMAGE_SIZE = (512, 512)
+DEFAULT_THRESHOLD = 0.35
+
+# Download model and metadata from Hugging Face Hub
 model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE, cache_dir=".")
 meta_path = hf_hub_download(repo_id=MODEL_REPO, filename=META_FILE, cache_dir=".")
+
+# Initialize ONNX Runtime session and load metadata
 session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
-metadata = json.load(open(meta_path, "r", encoding="utf-8"))
-# Preprocessing function (same as before)
+with open(meta_path, "r", encoding="utf-8") as f:
+    metadata = json.load(f)
+
+def escape_tag(tag: str) -> str:
+    """Escape underscores and parentheses for Markdown."""
+    return tag.replace("_", " ").replace("(", r"\\(").replace(")", r"\\)")
+
 def preprocess_image(pil_image: Image.Image) -> np.ndarray:
-    img = pil_image.convert("RGB").resize((512, 512))
+    """Convert image to RGB, resize, normalize, and rearrange dimensions."""
+    img = pil_image.convert("RGB").resize(IMAGE_SIZE)
     arr = np.array(img).astype(np.float32) / 255.0
     arr = np.transpose(arr, (2, 0, 1))
-    arr = np.expand_dims(arr, 0)
-    return arr
+    return np.expand_dims(arr, 0)
 
-# Inference function with output format option
-def tag_image(pil_image: Image.Image, output_format: str) -> str:
-    # Run model inference
+def run_inference(pil_image: Image.Image) -> np.ndarray:
+    """
+    Preprocess the image and run the ONNX model inference.
+    
+    Returns the refined logits as a numpy array.
+    """
     input_tensor = preprocess_image(pil_image)
     input_name = session.get_inputs()[0].name
-    initial_logits, refined_logits = session.run(None, {input_name: input_tensor})
+    # Only refined_logits are used (initial_logits is ignored)
+    _, 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):
+    """
+    Compute probabilities from logits and collect tag predictions.
+    
+    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).
+        all_artist_tags: All artist tags (with probabilities) regardless of threshold.
+    """
     probs = 1 / (1 + np.exp(-refined_logits))
-    probs = probs[0]
     idx_to_tag = metadata["idx_to_tag"]
     tag_to_category = metadata.get("tag_to_category", {})
     category_thresholds = metadata.get("category_thresholds", {})
-    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
+
+    results_by_cat = {}
+    prompt_tags_by_cat = {"artist": [], "character": [], "general": []}
+    all_artist_tags = []
+
     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 cat == "artist":
+            all_artist_tags.append((tag, float(prob)))
         if float(prob) >= thresh:
-            # add to category dictionary
             results_by_cat.setdefault(cat, []).append((tag, float(prob)))
-            if cat == 'artist':
-                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 cat in prompt_tags_by_cat:
+                prompt_tags_by_cat[cat].append((tag, float(prob)))
+    return results_by_cat, prompt_tags_by_cat, all_artist_tags
+
+def format_prompt_tags(prompt_tags_by_cat: dict, all_artist_tags: list) -> str:
+    """
+    Format the tags for prompt-style output.
+    
+    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)
+    
+    artist_tags = [escape_tag(tag) for tag, _ in prompt_tags_by_cat.get("artist", [])]
+    character_tags = [escape_tag(tag) for tag, _ in prompt_tags_by_cat.get("character", [])]
+    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
+    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))
+    return ", ".join(prompt_tags) if prompt_tags else "No tags predicted."
+
+def format_detailed_output(results_by_cat: dict, all_artist_tags: list) -> str:
+    """
+    Format the tags for detailed output.
+    
+    Returns a Markdown-formatted string listing tags by category.
+    """
+    if not results_by_cat:
+        return "No tags predicted for this image."
 
+    # Include an artist tag even if below threshold
+    if "artist" not in results_by_cat and all_artist_tags:
+        best_artist_tag, best_artist_prob = max(all_artist_tags, key=lambda item: item[1])
+        results_by_cat["artist"] = [(best_artist_tag, best_artist_prob)]
+    
+    lines = ["**Predicted Tags by Category:**  \n"]
+    for cat, tag_list in results_by_cat.items():
+        tag_list.sort(key=lambda x: x[1], reverse=True)
+        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
+    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."""
+    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)
+    
     if output_format == "Prompt-style Tags":
-        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("(", r"\\(").replace(")", r"\\)") for tag, prob in artist_tags_with_probs]
-        character_prompt_tags = [tag.replace("_", " ").replace("(", r"\\(").replace(")", r"\\)") for tag, prob in character_tags_with_probs]
-        general_prompt_tags = [tag.replace("_", " ").replace("(", r"\\(").replace(")", r"\\)") for tag, prob in general_tags_with_probs]
-
-        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("(", r"\\(").replace(")", r"\\)")] + prompt_tags
-
-
-        if not prompt_tags:
-            return "No tags predicted."
-        return ", ".join(prompt_tags)
-    else:  # Detailed output
-        if not results_by_cat:
-            return "No tags predicted for this image."
-
-        # Ensure artist tag in detailed output even if below threshold
-        if 'artist' not in results_by_cat and all_artist_tags_probs:
-            best_artist_tag, best_artist_prob = max(all_artist_tags_probs, key=lambda item: item[1])
-            results_by_cat['artist'] = [(best_artist_tag, best_artist_prob)]
-
-
-        lines = []
-        lines.append("**Predicted Tags by Category:**  \n")  # (Markdown newline: two spaces + newline)
-        for cat, tag_list in results_by_cat.items():
-            # sort tags in this category by probability descending
-            tag_list.sort(key=lambda x: x[1], reverse=True)
-            lines.append(f"**Category: {cat}** – {len(tag_list)} tags")
-            for tag, prob in tag_list:
-                tag_pretty = tag.replace("_", " ").replace("(", r"\\(").replace(")", r"\\)") # Escape parentheses here with raw string
-                lines.append(f"- {tag_pretty} (Prob: {prob:.3f})")
-            lines.append("")  # blank line between categories
-        return "\n".join(lines)
+        return format_prompt_tags(prompt_tags_by_cat, all_artist_tags)
+    else:
+        return format_detailed_output(results_by_cat, all_artist_tags)
 
 # Build the Gradio Blocks UI
-demo = gr.Blocks(theme="gradio/soft")  # using a built-in theme for nicer styling
+demo = gr.Blocks(theme="gradio/soft")
 
 with demo:
-    # Header Section
-    gr.Markdown("# 🏷️ Camie Tagger – Anime Image Tagging\nThis demo uses an ONNX model of Camie Tagger to label anime illustrations with tags. Upload an image 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. You can choose a concise prompt-style output or a detailed category-wise breakdown.)*")
-    # Input/Output Section
+    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."
+    )
+    gr.Markdown(
+        "*(Note: The model will predict 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():
-        # Left column: Image input and format selection
         with gr.Column():
             image_in = gr.Image(type="pil", label="Input Image")
-            format_choice = gr.Radio(choices=["Prompt-style Tags", "Detailed Output"], value="Prompt-style Tags", label="Output Format")
+            format_choice = gr.Radio(
+                choices=["Prompt-style Tags", "Detailed Output"],
+                value="Prompt-style Tags",
+                label="Output Format"
+            )
             tag_button = gr.Button("🔍 Tag Image")
-        # Right column: Output display
         with gr.Column():
-            output_box = gr.Markdown("")  # will display the result in Markdown (supports bold, lists, etc.)
-    # Link the button click to the function
+            output_box = gr.Markdown("")  # Markdown output for formatted results
+
     tag_button.click(fn=tag_image, inputs=[image_in, format_choice], outputs=output_box)
-    # Footer/Info
-    gr.Markdown("----\n**Model:** [Camie Tagger ONNX](https://huggingface.co/AngelBottomless/camie-tagger-onnxruntime)   •   **Base Model:** Camais03/camie-tagger (61% F1 on 70k tags)   •   **ONNX Runtime:** for efficient CPU inference​:contentReference[oaicite:6]{index=6}   •   *Demo built with Gradio Blocks.*")
+    
+    gr.Markdown(
+        "----\n"
+        "**Model:** [Camie Tagger ONNX](https://huggingface.co/AngelBottomless/camie-tagger-onnxruntime)   •   "
+        "**Base Model:** Camais03/camie-tagger (61% F1 on 70k tags)   •   **ONNX Runtime:** for efficient CPU inference   •   "
+        "*Demo built with Gradio Blocks.*"
+    )
 
-# Launch the app (automatically handled in Spaces)
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()