Update app.py

app.py

@@ -4,161 +4,121 @@ import numpy as np
 from PIL import Image
 import json
 from huggingface_hub import hf_hub_download
-from functools import partial # Import partial
 
-# --- Constants ---
+# Load model and metadata at startup (same as before)
 MODEL_REPO = "AngelBottomless/camie-tagger-onnxruntime"
 MODEL_FILE = "camie_tagger_initial.onnx"
 META_FILE = "metadata.json"
-IMAGE_SIZE = (512, 512)
-DEFAULT_THRESHOLD = 0.35
-
-
-# --- Helper Functions ---
-def download_model_and_metadata(repo_id: str, model_filename: str, meta_filename: str, cache_dir: str = "."):
-    """Downloads the ONNX model and metadata from Hugging Face Hub."""
-    model_path = hf_hub_download(repo_id=repo_id, filename=model_filename, cache_dir=cache_dir)
-    meta_path = hf_hub_download(repo_id=repo_id, filename=meta_filename, cache_dir=cache_dir)
-    return model_path, meta_path
-
-def load_model_session(model_path: str) -> ort.InferenceSession:
-    """Loads the ONNX model inference session."""
-    return ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
-
-def load_metadata(meta_path: str) -> dict:
-    """Loads the metadata from the JSON file."""
-    with open(meta_path, "r", encoding="utf-8") as f:
-        return json.load(f)
-
-def preprocess_image(pil_image: Image.Image, image_size: tuple = IMAGE_SIZE) -> np.ndarray:
-    """Preprocesses the PIL image to numpy array for model input."""
-    img = pil_image.convert("RGB").resize(image_size)
+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=".")
+session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
+metadata = json.load(open(meta_path, "r", encoding="utf-8"))
+# Preprocessing function (same as before)
+def preprocess_image(pil_image: Image.Image) -> np.ndarray:
+    img = pil_image.convert("RGB").resize((512, 512))
     arr = np.array(img).astype(np.float32) / 255.0
     arr = np.transpose(arr, (2, 0, 1))
     arr = np.expand_dims(arr, 0)
     return arr
 
-def apply_sigmoid(logits: np.ndarray) -> np.ndarray:
-    """Applies sigmoid function to logits to get probabilities."""
-    return 1 / (1 + np.exp(-logits))
-
-def extract_tags_from_probabilities(probs: np.ndarray, metadata: dict, threshold: float = DEFAULT_THRESHOLD) -> dict:
-    """Extracts tags and probabilities from the model output probabilities."""
+# Inference function with output format option
+def tag_image(pil_image: Image.Image, output_format: str) -> str:
+    # Run model inference
+    input_tensor = preprocess_image(pil_image)
+    input_name = session.get_inputs()[0].name
+    initial_logits, refined_logits = session.run(None, {input_name: input_tensor})
+    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", {})
-    results_by_cat = {}
-    all_artist_tags_probs = []
-
+    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)))
-        thresh = category_thresholds.get(cat, threshold)
+            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
             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)))
 
-    return results_by_cat, all_artist_tags_probs
-
-def format_prompt_style_output(results_by_cat: dict, all_artist_tags_probs: list) -> str:
-    """Formats the output as a comma-separated prompt-style string."""
-    artist_tags_with_probs = results_by_cat.get('artist', [])
-    character_tags_with_probs = results_by_cat.get('character', [])
-    general_tags_with_probs = results_by_cat.get('general', [])
-
-    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)
+    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_probs]
+        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
+        prompt_tags = artist_prompt_tags + character_prompt_tags + general_prompt_tags
 
-    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]) if all_artist_tags_probs else (None, None)
-        if best_artist_tag: # Check if best_artist_tag is not None
+        # 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)
 
-def format_detailed_output(results_by_cat: dict, all_artist_tags_probs: list) -> str:
-    """Formats the output as a detailed markdown string with categories and probabilities."""
-    if not results_by_cat:
-        return "No tags predicted for this image."
+        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."
 
-    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]) if all_artist_tags_probs else (None, None)
-        if best_artist_tag: # Check if best_artist_tag is not None
+        # 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")
-    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:
-            tag_pretty = tag.replace("_", " ").replace("(", r"\\(").replace(")", r"\\)")
-            lines.append(f"- {tag_pretty} (Prob: {prob:.3f})")
-        lines.append("")
-    return "\n".join(lines)
-
-# --- Inference Function ---
-def tag_image(pil_image: Image.Image, output_format: str, session: ort.InferenceSession, metadata: dict) -> str:
-    """Tags the image and formats the output based on the selected format."""
-    if pil_image is None:
-        return "Please upload an image."
-
-    input_tensor = preprocess_image(pil_image)
-    input_name = session.get_inputs()[0].name
-    initial_logits, refined_logits = session.run(None, {input_name: input_tensor})
-    probs = apply_sigmoid(refined_logits)[0]  # Apply sigmoid and get probabilities for the first (and only) image in batch
-
-    results_by_cat, all_artist_tags_probs = extract_tags_from_probabilities(probs, metadata)
-
-    if output_format == "Prompt-style Tags":
-        return format_prompt_style_output(results_by_cat, all_artist_tags_probs)
-    else:  # Detailed Output
-        return format_detailed_output(results_by_cat, all_artist_tags_probs)
-
-
-# --- Gradio UI ---
-def create_gradio_interface(session: ort.InferenceSession, metadata: dict) -> gr.Blocks:
-    """Creates the Gradio Blocks interface."""
-    demo = gr.Blocks(theme="gradio/soft")
-
-    with demo:
-        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.)*")
-
-        with gr.Row():
-            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")
-                tag_button = gr.Button("🔍 Tag Image")
-            with gr.Column():
-                output_box = gr.Markdown("")
-
-        # Create a partial function with session and metadata pre-filled
-        tag_image_with_model = partial(tag_image, session=session, metadata=metadata)
-
-        tag_button.click(
-            fn=tag_image_with_model, # Use the partially applied function
-            inputs=[image_in, format_choice],
-            outputs=output_box,
-        )
-
-        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.*")
-    return demo
 
-# --- Main Script ---
-if __name__ == "__main__":
-    model_path, meta_path = download_model_and_metadata(MODEL_REPO, MODEL_FILE, META_FILE)
-    session = load_model_session(model_path)
-    metadata = load_metadata(meta_path)
-    demo = create_gradio_interface(session, metadata)
-    demo.launch()
+        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)
+
+# Build the Gradio Blocks UI
+demo = gr.Blocks(theme="gradio/soft")  # using a built-in theme for nicer styling
+
+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
+    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")
+            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
+    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.*")
+
+# Launch the app (automatically handled in Spaces)
+demo.launch()