Method predict

modules/sam3_info_extractor.py:250–318

SAM3推理 Args: image_path: 图片路径 prompts: 提示词列表 score_threshold: 置信度阈值 min_area: 最小面积阈值 Returns: 元素列表

(self, image_path: str, prompts: List[str], 
                score_threshold: float = 0.5,
                min_area: int = 100)

        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:
            self._processor.reset_all_prompts(state)
            result_state = self._processor.set_text_prompt(prompt=prompt, state=state)
            
            masks = result_state.get("masks", [])
            boxes = result_state.get("boxes", [])
            scores = result_state.get("scores", [])
            
            num_masks = masks.shape[0] if (isinstance(masks, torch.Tensor) and masks.dim() > 0) else len(masks)
            
            for i in range(num_masks):
                score = scores[i]
                score_val = score.item() if hasattr(score, 'item') else float(score)
                
                if score_val < score_threshold:
                    continue
                
                # 提取bbox
                box = boxes[i]
                bbox = box.cpu().numpy().tolist() if isinstance(box, torch.Tensor) else box
                bbox = [int(coord) for coord in bbox]
                
                # 检查面积
                area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
                if area < min_area:
                    continue
                
                # 提取mask
                mask = masks[i]
                binary_mask = mask.cpu().numpy() if isinstance(mask, torch.Tensor) else np.array(mask)
                if binary_mask.ndim > 2:
                    binary_mask = binary_mask.squeeze()
                binary_mask = (binary_mask > 0.5).astype(np.uint8) * 255
                
                # 提取polygon
                polygon = self._extract_polygon(binary_mask, min_area)