Create app.py

app.py

@@ -0,0 +1,186 @@
+import os
+import gradio as gr
+import huggingface_hub
+import numpy as np
+import onnxruntime as rt
+import pandas as pd
+from PIL import Image
+
+# 模型配置
+MODEL_REPO = "SmilingWolf/wd-swinv2-tagger-v3"  # 默认模型
+MODEL_FILENAME = "model.onnx"
+LABEL_FILENAME = "selected_tags.csv"
+HF_TOKEN = os.environ.get("HF_TOKEN", "")
+
+# 标签处理配置
+kaomojis = [
+    "0_0",
+    "(o)_(o)",
+    "+_+",
+    "+_-",
+    "._.",
+    "<o>_<o>",
+    "<|>_<|>",
+    "=_=",
+    ">_<",
+    "3_3",
+    "6_9",
+    ">_o",
+    "@_@",
+    "^_^",
+    "o_o",
+    "u_u",
+    "x_x",
+    "|_|",
+    "||_||",
+]
+
+class Tagger:
+    def __init__(self):
+        self.model = None
+        self.tag_names = []
+        self.model_size = None
+        self._init_model()
+
+    def _init_model(self):
+        """初始化模型和标签"""
+        # 下载模型文件
+        label_path = huggingface_hub.hf_hub_download(
+            MODEL_REPO, 
+            LABEL_FILENAME,
+            token=HF_TOKEN
+        )
+        model_path = huggingface_hub.hf_hub_download(
+            MODEL_REPO,
+            MODEL_FILENAME,
+            token=HF_TOKEN
+        )
+        
+        # 加载标签
+        tags_df = pd.read_csv(label_path)
+        self.tag_names = tags_df["name"].tolist()
+        self.categories = {
+            "rating": np.where(tags_df["category"] == 9)[0],
+            "general": np.where(tags_df["category"] == 0)[0],
+            "character": np.where(tags_df["category"] == 4)[0]
+        }
+        
+        # 加载ONNX模型
+        self.model = rt.InferenceSession(model_path)
+        self.model_size = self.model.get_inputs()[0].shape[1]
+
+    def _preprocess(self, img):
+        """图像预处理"""
+        # 转换为RGB
+        if img.mode != "RGB":
+            img = img.convert("RGB")
+        
+        # 填充为正方形
+        size = max(img.size)
+        padded = Image.new("RGB", (size, size), (255, 255, 255))
+        padded.paste(img, ((size - img.width)//2, (size - img.height)//2))
+        
+        # 调整尺寸
+        if size != self.model_size:
+            padded = padded.resize((self.model_size, self.model_size), Image.BICUBIC)
+        
+        # 转换为BGR格式
+        return np.array(padded)[:, :, ::-1].astype(np.float32)
+
+    def predict(self, img, general_thresh=0.35, character_thresh=0.85):
+        """执行预测"""
+        # 预处理
+        img_data = self._preprocess(img)[np.newaxis]
+        
+        # 运行模型
+        input_name = self.model.get_inputs()[0].name
+        outputs = self.model.run(None, {input_name: img_data})[0][0]
+        
+        # 组织结果
+        results = {
+            "ratings": {},
+            "general": {},
+            "characters": {}
+        }
+        
+        # 处理评分标签
+        for idx in self.categories["rating"]:
+            tag = self.tag_names[idx].replace("_", " ")
+            results["ratings"][tag] = float(outputs[idx])
+        
+        # 处理通用标签
+        for idx in self.categories["general"]:
+            if outputs[idx] > general_thresh:
+                tag = self.tag_names[idx].replace("_", " ")
+                results["general"][tag] = float(outputs[idx])
+        
+        # 处理角色标签
+        for idx in self.categories["character"]:
+            if outputs[idx] > character_thresh:
+                tag = self.tag_names[idx].replace("_", " ")
+                results["characters"][tag] = float(outputs[idx])
+        
+        # 排序结果
+        results["general"] = dict(sorted(
+            results["general"].items(),
+            key=lambda x: x[1],
+            reverse=True
+        ))
+        
+        return results
+
+# 创建Gradio界面
+with gr.Blocks(theme=gr.themes.Soft(), title="AI图像标签分析器") as demo:
+    gr.Markdown("# 🖼️ AI图像标签分析器")
+    gr.Markdown("上传图片自动分析图像内容标签")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            img_input = gr.Image(type="pil", label="上传图片")
+            with gr.Accordion("高级设置", open=False):
+                general_slider = gr.Slider(0, 1, 0.35, 
+                    label="通用标签阈值", 
+                    info="值越高标签越少但更准确")
+                char_slider = gr.Slider(0, 1, 0.85,
+                    label="角色标签阈值",
+                    info="推荐保持较高阈值")
+            analyze_btn = gr.Button("开始分析", variant="primary")
+
+        with gr.Column(scale=2):
+            with gr.Tabs():
+                with gr.TabItem("🏷️ 通用标签"):
+                    general_tags = gr.Label(label="检测到的通用标签")
+                with gr.TabItem("👤 角色标签"):
+                    char_tags = gr.Label(label="检测到的角色标签")
+                with gr.TabItem("⭐ 评分标签"):
+                    rating_tags = gr.Label(label="图像评级标签")
+            
+            output_text = gr.Textbox(label="标签文本",
+                placeholder="生成的标签文本将显示在这里...")
+
+    # 处理逻辑
+    def process_image(img, gen_thresh, char_thresh):
+        tagger = Tagger()
+        results = tagger.predict(img, gen_thresh, char_thresh)
+        
+        # 格式化文本输出
+        tag_text = ", ".join(results["general"].keys())
+        if results["characters"]:
+            tag_text += ", " + ", ".join(results["characters"].keys())
+        
+        return {
+            general_tags: results["general"],
+            char_tags: results["characters"],
+            rating_tags: results["ratings"],
+            output_text: tag_text
+        }
+
+    analyze_btn.click(
+        process_image,
+        inputs=[img_input, general_slider, char_slider],
+        outputs=[general_tags, char_tags, rating_tags, output_text]
+    )
+
+# 启动应用
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)