JSar's picture
Create app.py
89683ee verified
Raw
History Blame Contribute Delete
616 Bytes
import gradio as gr
from transformers import pipeline
# Load the model
model_name = "SmilingWolf/wd-eva02-large-tagger-v3"
tagger = pipeline("image-classification", model=model_name)
# Define the prediction function
def predict(image):
results = tagger(image)
return {result["label"]: result["score"] for result in results}
# Create a Gradio interface
iface = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil"),
outputs=gr.Label(num_top_classes=10),
title="Danbooru AI Tagger",
description="Automatically tag Danbooru-style images using AI."
)
# Launch the interface
iface.launch()