Update app.py

app.py

@@ -1,115 +1,184 @@
 import gradio as gr
 import torch
 import timm
-from timm.data import resolve_data_config
-from timm.data.transforms_factory import create_transform
-from huggingface_hub import hf_hub_download
-import json
-import os
 from PIL import Image
+import json
+from torchvision import transforms
+from huggingface_hub import hf_hub_download
 
 # --- 配置 ---
-# 必须与您训练脚本中的 target_repo 一致
-REPO_ID = "telecomadm1145/convnext_dinov3_tagger_test_2w_asl_frozen"
-# 必须与您训练时的模型架构名称一致
-# 注意：如果在 Space 运行时报错找不到该模型名称，请尝试改为标准的 'convnext_base'
-MODEL_NAME = 'convnext_base.dinov3_lvd1689m' 
-TAGS_FILENAME = "tag_map.json"
+REPO_ID = "telecomadm1145/convnext_dinov3_tagger_test_epoch_4_asl_letterbox" 
 MODEL_FILENAME = "pytorch_model.bin"
-THRESHOLD = 0.25  # 显示标签的置信度阈值
+TAGS_FILENAME = "tag_map.json"
+MODEL_NAME = "convnext_base.dinov3_lvd1689m"
+INPUT_SIZE = (512, 512)
+
+# --- 1. 预处理 (Letterbox) ---
+class LetterboxPad:
+    def __init__(self, size, fill=(255, 255, 255)):
+        self.size = size if isinstance(size, tuple) else (size, size)
+        self.fill = fill
+
+    def __call__(self, img):
+        w, h = img.size
+        target_h, target_w = self.size
+        scale = min(target_w / w, target_h / h)
+        new_w = int(w * scale)
+        new_h = int(h * scale)
+        img = img.resize((new_w, new_h), Image.BICUBIC)
+        new_img = Image.new("RGB", (target_w, target_h), self.fill)
+        paste_x = (target_w - new_w) // 2
+        paste_y = (target_h - new_h) // 2
+        new_img.paste(img, (paste_x, paste_y))
+        return new_img
+
+def build_transform(size):
+    return transforms.Compose([
+        LetterboxPad(size),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
 
-print(f"Loading resources from {REPO_ID}...")
+# --- 2. 加载资源与分组 ---
+print("Loading model and tags...")
+device = torch.device("cpu")
+
+# 存储不同组的 (name, index) 列表
+tag_groups = {
+    "rating": [],
+    "character": [],
+    "general": []
+}
 
-# 1. 下载并加载标签映射
 try:
-    tags_path = hf_hub_download(repo_id=REPO_ID, filename=TAGS_FILENAME)
-    with open(tags_path, 'r') as f:
-        tag_map = json.load(f)
-    # 将 {'tag': 0, 'tag2': 1} 反转为 {0: 'tag', 1: 'tag2'} 以便索引
-    idx_to_tag = {v: k for k, v in tag_map.items()}
-    num_classes = len(tag_map)
-    print(f"Loaded {num_classes} tags.")
+    json_path = hf_hub_download(repo_id=REPO_ID, filename=TAGS_FILENAME)
+    with open(json_path, 'r') as f:
+        grouped_json = json.load(f)
+    
+    # 解析分组: 假设 JSON 结构为 {"rating": {"safe": 0, ...}, "general": ...}
+    total_tags = 0
+    for group_key, tags_dict in grouped_json.items():
+        # 兼容处理：确保 key 是我们预期的，如果只有 standard tags 可能会归类到 general
+        target_group = group_key if group_key in tag_groups else "general"
+        
+        for name, idx in tags_dict.items():
+            tag_groups[target_group].append((name, int(idx)))
+            total_tags += 1
+            
+    print(f"Loaded {total_tags} tags.")
+    print(f" - Rating: {len(tag_groups['rating'])}")
+    print(f" - Character: {len(tag_groups['character'])}")
+    print(f" - General: {len(tag_groups['general'])}")
+    
 except Exception as e:
     print(f"Error loading tags: {e}")
-    # 提供一个用于调试的空映射，防止 App 直接崩溃
-    idx_to_tag = {}
-    num_classes = 12476  
+    total_tags = 12000 # Fallback
 
-# 2. 准备模型
-print(f"Creating model: {MODEL_NAME}")
+# 加载模型
+model = timm.create_model(MODEL_NAME, pretrained=False, num_classes=total_tags)
 try:
-    # 创建模型结构
-    model = timm.create_model(MODEL_NAME, pretrained=False, num_classes=num_classes)
-    
-    # 下载并加载权重
     model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILENAME)
-    state_dict = torch.load(model_path, map_location='cpu')
+    state_dict = torch.load(model_path, map_location=device)
     model.load_state_dict(state_dict)
-    model.eval()
-    print("Model loaded successfully.")
+    print("Model weights loaded.")
 except Exception as e:
-    print(f"Error loading model: {e}")
-    print("Fallback: Attempting to use generic 'convnext_base'...")
-    try:
-        # 如果特定的 dinov3 命名在普通环境中不可用，尝试使用基础架构
-        model = timm.create_model('convnext_base', pretrained=False, num_classes=num_classes)
-        model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILENAME)
-        state_dict = torch.load(model_path, map_location='cpu')
-        model.load_state_dict(state_dict)
-        model.eval()
-        print("Fallback model loaded successfully.")
-    except Exception as e2:
-        raise RuntimeError(f"Failed to load model: {e2}")
+    print(f"Error loading weights: {e}")
 
-# 3. 准备图像预处理 (Transforms)
-config = resolve_data_config({}, model=model)
-transform = create_transform(**config)
+model.to(device)
+model.eval()
+transform = build_transform(INPUT_SIZE)
 
-# 4. 推理函数
+# --- 3. 推理逻辑 ---
 @torch.no_grad()
-def predict(image):
+def predict(image, threshold_gen, threshold_char):
     if image is None:
-        return {}
+        return {}, {}, {}
     
-    # 转换图像 (PIL -> Tensor, Normalize, Resize)
-    img_tensor = transform(image).unsqueeze(0) # Add batch dimension: [1, 3, H, W]
+    img_tensor = transform(image).unsqueeze(0).to(device)
+    logits = model(img_tensor)[0] # Shape: [num_classes]
     
-    # 前向传播
-    logits = model(img_tensor)
+    # --- A. 处理 Rating (Softmax) ---
+    rating_res = {}
+    if tag_groups["rating"]:
+        # 提取 rating 对应的 logits
+        r_indices = [idx for _, idx in tag_groups["rating"]]
+        r_names = [name for name, _ in tag_groups["rating"]]
+        
+        # 将 indices 转为 tensor 以便索引
+        r_indices_tensor = torch.tensor(r_indices, device=device)
+        r_logits = logits[r_indices_tensor]
+        
+        # 核心修改：对 Rating 组内进行 Softmax
+        r_probs = torch.nn.functional.softmax(r_logits, dim=0)
+        
+        for name, prob in zip(r_names, r_probs):
+            rating_res[name] = float(prob)
     
-    # 多标签分类使用 Sigmoid
-    probs = torch.sigmoid(logits)[0]
+    # --- B. 处理 Character (Sigmoid + Threshold) ---
+    char_res = {}
+    if tag_groups["character"]:
+        c_indices = [idx for _, idx in tag_groups["character"]]
+        c_names = [name for name, _ in tag_groups["character"]]
+        
+        c_indices_tensor = torch.tensor(c_indices, device=device)
+        c_logits = logits[c_indices_tensor]
+        c_probs = torch.sigmoid(c_logits) # 多标签使用 Sigmoid
+        
+        for name, prob in zip(c_names, c_probs):
+            if prob > threshold_char:
+                char_res[name] = float(prob)
+
+    # --- C. 处理 General (Sigmoid + Threshold) ---
+    gen_res = {}
+    if tag_groups["general"]:
+        g_indices = [idx for _, idx in tag_groups["general"]]
+        g_names = [name for name, _ in tag_groups["general"]]
+        
+        g_indices_tensor = torch.tensor(g_indices, device=device)
+        g_logits = logits[g_indices_tensor]
+        g_probs = torch.sigmoid(g_logits) # 多标签使用 Sigmoid
+        
+        for name, prob in zip(g_names, g_probs):
+            if prob > threshold_gen:
+                gen_res[name] = float(prob)
     
-    # 整理结果
-    results = {}
-    for idx, score in enumerate(probs):
-        score_val = float(score)
-        if score_val > THRESHOLD:
-            tag_name = idx_to_tag.get(idx, f"Unknown_{idx}")
-            results[tag_name] = score_val
-            
-    # 按照置信度排序 (Gradio 的 Label 组件会自动排序，但手动排序方便调试)
-    # 限制返回前 50 个标签，防止 UI 过于拥挤
-    sorted_results = dict(sorted(results.items(), key=lambda x: x[1], reverse=True)[:50])
+    # 排序
+    rating_res = dict(sorted(rating_res.items(), key=lambda x: x[1], reverse=True))
+    char_res = dict(sorted(char_res.items(), key=lambda x: x[1], reverse=True))
+    gen_res = dict(sorted(gen_res.items(), key=lambda x: x[1], reverse=True))
     
-    return sorted_results
+    return rating_res, char_res, gen_res
 
-# 5. 构建 Gradio 界面
+# --- 4. 界面 ---
 with gr.Blocks() as demo:
+    gr.Markdown(f"# Anime Tagger (DINOv3)\nModel: {REPO_ID}")
+    
     with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=1):
             input_img = gr.Image(type="pil", label="Input Image")
-            run_btn = gr.Button("Run Tagger", variant="primary")
-        
-        with gr.Column():
-            # Label 组件非常适合显示分类概率
-            output_tags = gr.Label(label="Detected Tags", num_top_classes=30)
-    
-    run_btn.click(fn=predict, inputs=input_img, outputs=output_tags)
-    
-    examples = []
-    if examples:
-        gr.Examples(examples=examples, inputs=input_img)
+            run_btn = gr.Button("Tag It!", variant="primary")
+            
+            gr.Markdown("### Thresholds")
+            # 为不同类别设置不同的阈值通常更好，Character 往往需要更低的阈值来召回
+            threshold_gen = gr.Slider(0.0, 1.0, value=0.25, step=0.05, label="General Tags Threshold")
+            threshold_char = gr.Slider(0.0, 1.0, value=0.15, step=0.05, label="Character Threshold")
+
+        with gr.Column(scale=1):
+            # 分开显示
+            gr.Markdown("### 1. Rating (Softmax)")
+            out_rating = gr.Label(label="Rating")
+            
+            gr.Markdown("### 2. Characters")
+            out_char = gr.Label(label="Characters", num_top_classes=10)
+            
+            gr.Markdown("### 3. General Tags")
+            out_gen = gr.Label(label="General Tags", num_top_classes=50)
+
+    run_btn.click(
+        fn=predict, 
+        inputs=[input_img, threshold_gen, threshold_char], 
+        outputs=[out_rating, out_char, out_gen]
+    )
 
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.launch()