app.py

import gradio as gr
import onnxruntime as ort
import numpy as np
from PIL import Image
import json
from huggingface_hub import hf_hub_download

# 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"])
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:
    """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))
    return np.expand_dims(arr, 0)

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
    # 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, 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: 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))
    idx_to_tag = metadata["idx_to_tag"]
    tag_to_category = metadata.get("tag_to_category", {})
    category_thresholds = metadata.get("category_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")
        # 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:
            results_by_cat.setdefault(cat, []).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.
    """
    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 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))
    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("") 
    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.
    """
    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)
    
    if output_format == "Prompt-style Tags":
        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")

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."
    )
    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"
            )
            # 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("")

    tag_button.click(
        fn=tag_image,
        inputs=[image_in, format_choice, threshold_slider],
        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.*"
    )

if __name__ == "__main__":
    demo.launch()