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
import torchvision.transforms as transforms

# Constants
MODEL_REPO = "AngelBottomless/camie-tagger-onnxruntime"
MODEL_FILE = "camie_tagger_initial.onnx"
META_FILE = "metadata.json"
DEFAULT_THRESHOLD = 0.32626262307167053  # 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=".")
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:
    """Process an image for inference using same preprocessing as training"""

    image_size=512

    # Initialize the same transform used during training
    transform = transforms.Compose([
        transforms.ToTensor(),
    ])

    img = pil_image # Use the PIL image directly

    # Convert RGBA or Palette images to RGB
    if img.mode in ('RGBA', 'P'):
        img = img.convert('RGB')

    # Get original dimensions
    width, height = img.size
    aspect_ratio = width / height

    # Calculate new dimensions to maintain aspect ratio
    if aspect_ratio > 1:
        new_width = image_size
        new_height = int(new_width / aspect_ratio)
    else:
        new_height = image_size
        new_width = int(new_height * aspect_ratio)

    # Resize with LANCZOS filter
    img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)

    # Create new image with padding
    new_image = Image.new('RGB', (image_size, image_size), (0, 0, 0))
    paste_x = (image_size - new_width) // 2
    paste_y = (image_size - new_height) // 2
    new_image.paste(img, (paste_x, paste_y))

    # Apply transforms (without normalization)
    img_tensor = transform(new_image)
    return img_tensor.numpy() # Convert the PyTorch tensor to NumPy array


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
    # Expand dimensions to make it (1, C, H, W)
    input_tensor_expanded = np.expand_dims(input_tensor, axis=0)
    # Only refined_logits are used (initial_logits is ignored)
    _, refined_logits = session.run(None, {input_name: input_tensor_expanded})
    return refined_logits[0]

def get_tags(refined_logits: np.ndarray, metadata: dict, default_threshold: float):
    """
    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: Dictionary for prompt-style output (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 = {}
    # For prompt style, only include character and general tags (artists handled separately)
    prompt_tags_by_cat = {"character": [], "general": []}
    all_artist_tags = []

    for idx, prob in enumerate(probs):
        tag = idx_to_tag[str(idx)]
        cat = tag_to_category.get(tag, "unknown")
        thresh = 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.
    Only the top artist tag is shown (regardless of threshold), and all character and general tags are shown.

    Returns a comma-separated string of escaped tags.
    """
    # Always select the best artist tag from all_artist_tags, regardless of threshold.
    best_artist_tag = None
    if all_artist_tags:
        best_artist = max(all_artist_tags, key=lambda item: item[1])
        best_artist_tag = escape_tag(best_artist[0])

    # Sort character and general tags by probability (descending)
    for cat in prompt_tags_by_cat:
        prompt_tags_by_cat[cat].sort(key=lambda x: x[1], reverse=True)

    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 = []
    if best_artist_tag:
        prompt_tags.append(best_artist_tag)
    prompt_tags.extend(character_tags)
    prompt_tags.extend(general_tags)

    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, threshold: float) -> str:
    """
    Run inference on the image and return formatted tags based on the chosen output format.

    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, default_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, adjust the threshold, and click **Tag Image** to see predictions."
    )
    gr.Markdown(
        "*(Note: In prompt-style output, only the top artist tag is displayed along with all character and general tags.)*"
    )
    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 default threshold value used in inference.
            threshold_slider = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                step=0.05,
                value=DEFAULT_THRESHOLD,
                label="Threshold"
            )
            tag_button = gr.Button("🔍 Tag Image")
        with gr.Column():
            output_box = gr.Markdown("")  # Markdown output for formatted results

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