Class SAM3Model

modules/sam3_info_extractor.py:214–365

SAM3模型封装

# ======================== SAM3模型封装 ========================
class SAM3Model(ModelWrapper):
    """SAM3模型封装"""
    
    def __init__(self, checkpoint_path: str, bpe_path: str, device: str = None):
        super().__init__()
        self.checkpoint_path = checkpoint_path
        self.bpe_path = bpe_path
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self._processor = None
        
        # 图像状态缓存
        self._state_cache = OrderedDict()
        self._max_cache_size = 3
        self._cache_lock = threading.Lock()
    
    def load(self):
        """加载SAM3模型"""
        if self._is_loaded:
            return
            
        print(f"[SAM3Model] 加载模型中... (设备: {self.device})")
        
        from sam3.model_builder import build_sam3_image_model
        from sam3.model.sam3_image_processor import Sam3Processor
        
        self._model = build_sam3_image_model(
            bpe_path=self.bpe_path,
            checkpoint_path=self.checkpoint_path,
            load_from_HF=False,
            device=self.device
        )
        self._processor = Sam3Processor(self._model)
        self._is_loaded = True
        
        print("[SAM3Model] 模型加载完成！")
    
    def predict(self, image_path: str, prompts: List[str], 
                score_threshold: float = 0.5,
                min_area: int = 100) -> List[Dict[str, Any]]:
        """
        SAM3推理
        
        Args:
            image_path: 图片路径
            prompts: 提示词列表
            score_threshold: 置信度阈值
            min_area: 最小面积阈值
            
        Returns:
            元素列表
        """
        if not self._is_loaded:
            self.load()
        
        state, pil_image = self._get_image_state(image_path)
        
        results = []
        for prompt in prompts: