Spaces:
Sleeping
Sleeping
| 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 | |
| # 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" | |
| 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 | |
| # 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", {}) | |
| default_threshold = 0.325 | |
| results_by_cat = {} # to store tags per category (for verbose output) | |
| prompt_tags = [] # to store tags for prompt-style output | |
| # Collect tags above thresholds | |
| 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 float(prob) >= thresh: | |
| # add to category dictionary | |
| results_by_cat.setdefault(cat, []).append((tag, float(prob))) | |
| # add to prompt list | |
| prompt_tags.append(tag.replace("_", " ")) | |
| if output_format == "Prompt-style Tags": | |
| if not prompt_tags: | |
| return "No tags predicted." | |
| # Join tags with commas (sorted by probability for relevance) | |
| # Sort prompt_tags by probability from results_by_cat (for better prompts ordering) | |
| prompt_tags.sort(key=lambda t: max([p for (tg, p) in results_by_cat[tag_to_category.get(t.replace(' ', '_'), 'unknown')] if tg == t.replace(' ', '_')]), reverse=True) | |
| return ", ".join(prompt_tags) | |
| else: # Detailed output | |
| if not results_by_cat: | |
| return "No tags predicted for this image." | |
| 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("_", " ") | |
| 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=gr.themes.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.) | |
| # Example images (if available in the repo) | |
| gr.Examples( | |
| examples=[["example1.jpg"], ["example2.png"]], # Example file paths (ensure these exist in the Space) | |
| inputs=image_in, | |
| outputs=output_box, | |
| fn=tag_image, | |
| cache_examples=True | |
| ) | |
| # 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() | |