Update app.py

app.py

@@ -1,141 +1,21 @@
-import subprocess
-import sys
-import os
-import re
-import glob
-
-# ============================================================
-# ✅ 终极修复：在 import transformers 之前，
-#    直接修改已安装的 transformers 包中的源码文件，
-#    让 initializer_range 字段同时接受 int 和 float。
-#    这样即使 ZeroGPU 重新反序列化模型也不会报错。
-# ============================================================
-
-def patch_transformers_source():
-    """
-    找到 transformers 包中所有 Sam3 相关的 configuration 文件，
-    把 initializer_range 的类型注解从严格的 int 或 float 改为兼容的 Union 类型，
-    或者直接移除类型校验。
-    """
-    import transformers
-    pkg_dir = os.path.dirname(transformers.__file__)
-    
-    # 找到所有可能的配置文件
-    patterns = [
-        os.path.join(pkg_dir, "models", "sam3", "*.py"),
-        os.path.join(pkg_dir, "models", "sam3", "**", "*.py"),
-    ]
-    
-    files_to_check = []
-    for pattern in patterns:
-        files_to_check.extend(glob.glob(pattern, recursive=True))
-    
-    if not files_to_check:
-        print(f"⚠️ 未找到 sam3 模型文件，尝试搜索整个 transformers 目录...")
-        # 搜索所有包含 Sam3 和 initializer_range 的文件
-        result = subprocess.run(
-            ["grep", "-rl", "initializer_range", os.path.join(pkg_dir, "models")],
-            capture_output=True, text=True
-        )
-        if result.stdout:
-            all_files = result.stdout.strip().split("\n")
-            # 只处理 sam3 相关的
-            files_to_check = [f for f in all_files if "sam3" in f.lower() or "sam_3" in f.lower()]
-            if not files_to_check:
-                # 如果没找到 sam3 特定的，搜索 configuration 文件
-                files_to_check = [f for f in all_files if "configuration" in f.lower()]
-    
-    patched_count = 0
-    for filepath in files_to_check:
-        try:
-            with open(filepath, "r") as f:
-                content = f.read()
-            
-            if "initializer_range" not in content:
-                continue
-            
-            original = content
-            
-            # 策略1: 把 initializer_range: int 改为 initializer_range: float
-            content = re.sub(
-                r'(initializer_range\s*:\s*)int(\s*=)',
-                r'\1float\2',
-                content
-            )
-            
-            # 策略2: 把 initializer_range: int = 0 改为 initializer_range: float = 0.0
-            content = re.sub(
-                r'(initializer_range\s*:\s*\w+\s*=\s*)(\d+)(\s*[,\n\)])',
-                lambda m: f'{m.group(1)}{float(int(m.group(2)))}{m.group(3)}',
-                content
-            )
-            
-            # 策略3: 如果有 validator 或 field_validator 针对 initializer_range 的严格类型检查
-            # 注释掉相关校验行
-            
-            if content != original:
-                with open(filepath, "w") as f:
-                    f.write(content)
-                patched_count += 1
-                print(f"✅ 已修补文件: {filepath}")
-            
-        except Exception as e:
-            print(f"⚠️ 处理文件 {filepath} 时出错: {e}")
-    
-    if patched_count == 0:
-        print("⚠️ 未找到需要修补的文件，尝试通用方案...")
-        # 通用方案：patch PretrainedConfig 的 __init_subclass__
-        patch_config_base_class()
-    else:
-        print(f"✅ 共修补了 {patched_count} 个文件")
-
-
-def patch_config_base_class():
-    """
-    如果找不到具体文件可改，就 patch PretrainedConfig 基类，
-    让所有配置类在实例化时自动容忍 int/float 互转。
-    """
-    from transformers import PretrainedConfig
-    
-    original_init_subclass = PretrainedConfig.__init_subclass__
-    original_init = PretrainedConfig.__init__
-    
-    # patch __setattr__ 让赋值时自动兼容
-    original_setattr = PretrainedConfig.__setattr__ if hasattr(PretrainedConfig, '__setattr__') else object.__setattr__
-    
-    def tolerant_setattr(self, name, value):
-        # 对 initializer_range 不做严格类型检查
-        if name == "initializer_range":
-            # 直接写入，跳过任何校验
-            self.__dict__[name] = value
-            return
-        try:
-            original_setattr(self, name, value)
-        except TypeError:
-            self.__dict__[name] = value
-    
-    PretrainedConfig.__setattr__ = tolerant_setattr
-
-
-# === 执行修复 ===
-print("🔧 正在修补 transformers 源码中的 initializer_range 类型问题...")
-patch_transformers_source()
-print("🔧 修补完成！")
-
-# === 现在才导入其他模块 ===
 import gradio as gr
 import torch
 import numpy as np
 import matplotlib.pyplot as plt
+import matplotlib
 from PIL import Image
 import gc
-import spaces 
+import os
+import spaces
 import cv2
-from transformers import Sam3Model, Sam3Processor
+from transformers import (
+    Sam3Model, Sam3Processor,
+    Sam3TrackerModel, Sam3TrackerProcessor,
+)
 
 HF_TOKEN = os.getenv("HF_TOKEN")
 MODELS = {}
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "cuda"
 
 def cleanup_memory():
     if MODELS:
@@ -143,156 +23,174 @@ def cleanup_memory():
     gc.collect()
     torch.cuda.empty_cache()
 
-def get_model():
-    model_id = "facebook/sam3"
-    if model_id in MODELS:
-        return MODELS[model_id]
+def get_model(model_type):
+    """
+    按需加载不同模型：
+    - sam3_image_text: 文本检测用 Sam3Model + Sam3Processor
+    - sam3_image_tracker: 样本/Box检测用 Sam3TrackerModel + Sam3TrackerProcessor
+    """
+    if model_type in MODELS:
+        return MODELS[model_type]
     
     cleanup_memory()
-    print("⏳ 正在加载 SAM 3 模型...")
-    
-    model = Sam3Model.from_pretrained(
-        model_id, 
-        token=HF_TOKEN, 
-        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-    ).to(device)
-    
-    processor = Sam3Processor.from_pretrained(
-        model_id, 
-        token=HF_TOKEN
-    )
+    print(f"⏳ 正在加载 {model_type} 模型...")
+    
+    if model_type == "sam3_image_text":
+        model = Sam3Model.from_pretrained("facebook/sam3", token=HF_TOKEN).to(device)
+        processor = Sam3Processor.from_pretrained("facebook/sam3", token=HF_TOKEN)
+    elif model_type == "sam3_image_tracker":
+        model = Sam3TrackerModel.from_pretrained("facebook/sam3", token=HF_TOKEN).to(device)
+        processor = Sam3TrackerProcessor.from_pretrained("facebook/sam3", token=HF_TOKEN)
+    else:
+        raise ValueError(f"未知模型类型: {model_type}")
     
-    MODELS[model_id] = (model, processor)
-    return MODELS[model_id]
+    MODELS[model_type] = (model, processor)
+    print(f"✅ {model_type} 加载完成。")
+    return MODELS[model_type]
 
 def overlay_masks(image, masks, alpha=0.6):
-    if image is None: 
+    if image is None:
         return None
-    if isinstance(image, np.ndarray): 
+    if isinstance(image, np.ndarray):
         image = Image.fromarray(image)
     image = image.convert("RGBA")
-    
-    if masks is None or len(masks) == 0: 
+
+    if masks is None or len(masks) == 0:
         return image.convert("RGB")
-        
-    if isinstance(masks, torch.Tensor): 
+
+    if isinstance(masks, torch.Tensor):
         masks = masks.cpu().numpy()
-        
+
     masks = masks.astype(np.uint8)
     if masks.ndim == 4: masks = masks[0]
     if masks.ndim == 3 and masks.shape[0] == 1: masks = masks[0]
     if masks.ndim == 2: masks = [masks]
-    
+
     n_masks = len(masks)
-    cmap = plt.get_cmap("rainbow", max(n_masks, 1))
+    try:
+        cmap = matplotlib.colormaps["rainbow"].resampled(max(n_masks, 1))
+    except AttributeError:
+        cmap = plt.get_cmap("rainbow", max(n_masks, 1))
+
     overlay_layer = Image.new("RGBA", image.size, (0, 0, 0, 0))
-    
     for i, mask in enumerate(masks):
         mask_img = Image.fromarray((mask * 255).astype(np.uint8))
-        if mask_img.size != image.size: 
+        if mask_img.size != image.size:
             mask_img = mask_img.resize(image.size, resample=Image.NEAREST)
-        
         rgb = [int(x * 255) for x in cmap(i)[:3]]
         color_layer = Image.new("RGBA", image.size, tuple(rgb) + (0,))
         mask_alpha = mask_img.point(lambda v: int(v * alpha) if v > 0 else 0)
         color_layer.putalpha(mask_alpha)
         overlay_layer = Image.alpha_composite(overlay_layer, color_layer)
-        
+
     return Image.alpha_composite(image, overlay_layer).convert("RGB")
 
+
+# ========== 文本描述检测 ==========
 @spaces.GPU
 def process_text_detection(image, text_query, threshold):
-    if not image or not text_query: 
+    if not image or not text_query:
         return None, "请输入图像和描述词"
-        
     try:
-        model, processor = get_model()
-        
+        model, processor = get_model("sam3_image_text")
+
         inputs = processor(
-            images=image, 
-            text=text_query, 
+            images=image,
+            text=text_query,
             return_tensors="pt"
         ).to(device)
-        
-        with torch.no_grad(): 
+
+        with torch.no_grad():
             outputs = model(**inputs)
-            
+
         results = processor.post_process_instance_segmentation(
-            outputs, 
-            threshold=threshold, 
-            mask_threshold=0.5, 
+            outputs,
+            threshold=threshold,
+            mask_threshold=0.5,
             target_sizes=inputs.get("original_sizes").tolist()
         )[0]
-        
+
         masks = results["masks"]
         result_img = overlay_masks(image, masks)
-        
+
         if len(masks) > 0:
             status = f"✅ 文本检测完成！找到 {len(masks)} 个目标。"
         else:
             status = "❓ 未找到目标，请调低阈值。"
-            
         return result_img, status
-        
+
     except Exception as e:
         return image, f"❌ 错误: {str(e)}"
 
+
+# ========== 样���截图检测 ==========
 @spaces.GPU
 def process_sample_detection(main_image, sample_image):
-    if not main_image or not sample_image: 
+    if not main_image or not sample_image:
         return None, "请上传主图和样本截图"
-        
     try:
-        model, processor = get_model()
-        
+        # ✅ 关键：box prompt 必须用 Sam3TrackerModel，不能用 Sam3Model
+        model, processor = get_model("sam3_image_tracker")
+
+        # Step 1: OpenCV 模板匹配，定位样本在主图中的位置
         main_cv = cv2.cvtColor(np.array(main_image), cv2.COLOR_RGB2BGR)
         sample_cv = cv2.cvtColor(np.array(sample_image), cv2.COLOR_RGB2BGR)
-        
+
         if sample_cv.shape[0] > main_cv.shape[0] or sample_cv.shape[1] > main_cv.shape[1]:
             return main_image, "❌ 错误：样本截图不能比主图还大！"
-            
+
         result = cv2.matchTemplate(main_cv, sample_cv, cv2.TM_CCOEFF_NORMED)
         min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
-        
+
         if max_val < 0.4:
             return main_image, f"❓ 未在主图中找到该样本 (最高匹配度: {max_val:.2f})。"
-            
+
         h, w = sample_cv.shape[:2]
-        
         box = [
-            max_loc[0], 
-            max_loc[1], 
-            max_loc[0] + w, 
+            max_loc[0],
+            max_loc[1],
+            max_loc[0] + w,
             max_loc[1] + h
         ]
-        
+
+        # Step 2: 用 Sam3TrackerProcessor 的 box prompt 做精细分割
         inputs = processor(
-            images=main_image, 
-            input_boxes=[[[box]]], 
+            images=main_image,
+            input_boxes=[[[box]]],
             return_tensors="pt"
         ).to(device)
-        
-        with torch.no_grad(): 
+
+        with torch.no_grad():
             outputs = model(**inputs)
-        
-        results = processor.post_process_instance_segmentation(
-            outputs, 
-            threshold=0.1, 
-            mask_threshold=0.5, 
-            target_sizes=inputs.get("original_sizes").tolist()
+
+        # Sam3Tracker 用 post_process_masks 而不是 post_process_instance_segmentation
+        masks = processor.post_process_masks(
+            outputs.pred_masks.cpu(),
+            inputs["original_sizes"],
+            binarize=True
         )[0]
-            
-        masks = results["masks"]
+
+        # masks 的形状是 [num_objects, num_masks, H, W]，取第一个物体的最佳 mask
+        if masks.ndim == 4:
+            # 取 IoU 最高的 mask
+            if hasattr(outputs, 'iou_scores') and outputs.iou_scores is not None:
+                scores = outputs.iou_scores.cpu().numpy()[0, 0]
+                best_idx = np.argmax(scores)
+                masks = masks[0, best_idx:best_idx+1]
+            else:
+                masks = masks[0, 0:1]
+
         result_img = overlay_masks(main_image, masks)
-        
         return result_img, f"✅ 样本检测成功！(匹配度: {max_val:.2f})"
-        
+
     except Exception as e:
         return main_image, f"❌ 错误: {str(e)}"
 
+
+# ========== Gradio 界面 ==========
 with gr.Blocks() as demo:
     gr.Markdown("# 🚀 SAM 3 自动检测工具 (双模式)")
-    
+
     with gr.Tabs():
         with gr.Tab("📝 文本描述检测"):
             with gr.Row():
@@ -305,8 +203,8 @@ with gr.Blocks() as demo:
                     t_img_out = gr.Image(type="pil", label="检测结果")
                     t_info = gr.Textbox(label="状态信息")
             t_btn.click(
-                process_text_detection, 
-                [t_img_in, t_query, t_thresh], 
+                process_text_detection,
+                [t_img_in, t_query, t_thresh],
                 [t_img_out, t_info]
             )
 
@@ -321,8 +219,8 @@ with gr.Blocks() as demo:
                     s_img_out = gr.Image(type="pil", label="检测结果")
                     s_info = gr.Textbox(label="状态信息")
             s_btn.click(
-                process_sample_detection, 
-                [s_img_main, s_img_sample], 
+                process_sample_detection,
+                [s_img_main, s_img_sample],
                 [s_img_out, s_info]
             )