Spaces:
Running
Running
| 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 the ONNX model and metadata once at startup (optimizes performance) | |
| MODEL_REPO = "AngelBottomless/camie-tagger-onnxruntime" | |
| MODEL_FILE = "camie_tagger_initial.onnx" # using the smaller initial model for speed | |
| META_FILE = "metadata.json" | |
| # Download model and metadata from HF Hub (cache_dir="." will cache in the Space) | |
| 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: resize image to 512x512 and normalize to match training | |
| 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 # scale pixel values to [0,1] | |
| arr = np.transpose(arr, (2, 0, 1)) # HWC -> CHW | |
| arr = np.expand_dims(arr, 0) # add batch dimension -> (1,3,512,512) | |
| return arr | |
| # Inference: run the ONNX model and collect tags above threshold | |
| def predict_tags(pil_image: Image.Image) -> str: | |
| # 1. Preprocess image to numpy | |
| input_tensor = preprocess_image(pil_image) | |
| # 2. Run model (both initial and refined logits are output) | |
| input_name = session.get_inputs()[0].name | |
| initial_logits, refined_logits = session.run(None, {input_name: input_tensor}) | |
| # 3. Convert logits to probabilities (using sigmoid since multi-label) | |
| probs = 1 / (1 + np.exp(-refined_logits)) # shape (1, 70527) | |
| probs = probs[0] # remove batch dim -> (70527,) | |
| # 4. Thresholding: get tag names for which probability >= category threshold (or default) | |
| idx_to_tag = metadata["idx_to_tag"] # map index -> tag string | |
| tag_to_category = metadata.get("tag_to_category", {}) # map tag -> category | |
| category_thresholds = metadata.get("category_thresholds", {})# category-specific thresholds | |
| default_threshold = 0.325 | |
| predicted_tags = [] | |
| for idx, prob in enumerate(probs): | |
| tag = idx_to_tag[str(idx)] | |
| cat = tag_to_category.get(tag, "unknown") | |
| threshold = category_thresholds.get(cat, default_threshold) | |
| if prob >= threshold: | |
| # Include this tag; replace underscores with spaces for readability | |
| predicted_tags.append(tag.replace("_", " ")) | |
| # 5. Return tags as comma-separated string | |
| if not predicted_tags: | |
| return "No tags found." | |
| # Join tags, maybe sorted by name or leave unsorted. Here we sort alphabetically for consistency. | |
| predicted_tags.sort() | |
| return ", ".join(predicted_tags) | |
| # Create a simple Gradio interface | |
| demo = gr.Interface( | |
| fn=predict_tags, | |
| inputs=gr.Image(type="pil", label="Upload Image"), | |
| outputs=gr.Textbox(label="Predicted Tags", lines=3), | |
| title="Camie Tagger (ONNX) – Simple Demo", | |
| description="Upload an anime/manga illustration to get relevant tags predicted by the Camie Tagger model.", | |
| # You can optionally add example images if available in the Space directory: | |
| examples=[["example1.jpg"], ["example2.png"]] # (filenames should exist in the Space) | |
| ) | |
| # Launch the app (in HF Spaces, just calling demo.launch() is typically not required; the Space will run app automatically) | |
| demo.launch() | |