Update app.py

app.py

@@ -5,120 +5,157 @@ 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"
-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))
+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)
     arr = np.array(img).astype(np.float32) / 255.0
     arr = np.transpose(arr, (2, 0, 1))
     arr = np.expand_dims(arr, 0)
     return arr
 
-# 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]
+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."""
     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 = {}
+    all_artist_tags_probs = []
+
     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)
+            all_artist_tags_probs.append((tag, float(prob)))
+        thresh = category_thresholds.get(cat, 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)))
 
-    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)
+    return results_by_cat, all_artist_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]
+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', [])
 
-        prompt_tags = artist_prompt_tags + character_prompt_tags + general_prompt_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)
 
-        # 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])
+    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
+
+    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
             prompt_tags = [best_artist_tag.replace("_", " ").replace("(", r"\\(").replace(")", r"\\)")] + prompt_tags
 
+    if not prompt_tags:
+        return "No tags predicted."
+    return ", ".join(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."
+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."
 
-        # 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])
+    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
             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("")
+
+        tag_button.click(
+            fn=tag_image,
+            inputs=[image_in, format_choice],
+            outputs=output_box,
+            extra_args=[session, metadata] # Pass session and metadata as extra arguments
+        )
+
+        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
 
-        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()
+# --- 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()