基于SAM-3的图像自动标注

本文介绍如何使用SAM-3构建一个图像自动标注工具,并利用这个工具创建了一个牙齿数据集。

基于SAM-3的图像自动标注

我最近遇到了一个任务:分割牙科图像中的单个牙齿,并为每颗牙齿分配 FDI 编号。

传统的做法是,使用标注工具逐张图像、逐颗牙齿地绘制多边形,从而构建一个标注数据集。由于有数百张图像,其中许多图像包含完整的 32 颗牙齿,这项工作耗时耗力。

在本文中,我将展示如何构建一个基于 SAM-3 的 Python 封装器。该封装器不仅为 SAM-3 提供了一个简单的入口点,用于分割图像中的任何对象,而且还创建了一个标注数据集,只需在标注工具中进行少量调整即可。这种方法可以应用于任何图像标注任务。

1、引言

计算机视觉中的分割是指将图像分割成有意义的区域,并根据每个像素所代表的含义为其分配标签的过程。图像分割并非将图像视为一个整体,而是将其分解,以便系统能够精确地识别哪些像素属于哪个对象或区域。

例如,在街景图像中,图像分割可以分离汽车、公共汽车、道路、建筑物和行人的像素,勾勒出每个对象的形状和边界。在医学成像中,图像分割可用于从牙科 X 光片中分离出单个牙齿,或从扫描图像中分离出心脏或肺部等器官,从而实现精确的测量和分析。

医学图像中肺部的分割

图像分割一直是计算机视觉领域最具挑战性的任务之一,尤其是在分割形状和大小各异的不规则物体时。

我最近遇到了一个问题:如何在牙科图像中分割人牙,并为每个分割后的图像分配一个 FDI 值。

对于这类任务,需要训练一个能够精确分割所有牙齿的模型。在此之前,我们需要先使用标注工具准备一个带标签的数据集。这需要手动在所有图像中绘制多边形来标记所有牙齿——这是一项非常耗时的任务,尤其是在每张图像中需要标注多颗牙齿,且数据集包含成百上千张图像的情况下。

我决定尝试使用 Meta 的 Segment Anything (SAM-3) 模型,希望能够预先准备标注数据,也就是说,SAM-3 模型能够分割出一些牙齿,然后可以通过图像标注工具进一步调整/改进。至少这样可以减少一些手动工作。

我找不到任何能够使用 SAM-3 分割批量图像并自动生成标注数据(例如最常见的 Common Objects in Context (COCO) 格式)的工具。因此,我开发了一个 SAM-3 的 Python 封装器,它抽象了大部分技术/实现细节,并提供了一个可导入的软件包,用于使用 SAM-3 模型进行实际分割和/或准备预标注数据。

结果超出了我的预期。

在本文中,我将带您了解我开发的 SAM-3 Python 封装库的开发过程,该封装库旨在使文本提示分割能够应用于实际场景。您将学习如何同时检测多种对象类型、以标注平台支持的格式导出标注,以及构建完整的数据集创建流程。

无论您是构建训练数据集、开发对象检测系统原型,还是仅仅想探索现代分割模型的可能性,本指南都将为您指明方向。

完整的代码可在 GitHub 代码库中找到。

让我们深入了解它的工作原理。

2、Segment Anything Model (SAM-3) 模型

SAM 是一种计算机视觉基础模型,旨在执行提示式视觉分割,即根据用户的提示将图像或视频分割成有意义的片段。

在传统的分割系统中,模型通常需要针对特定​​对象类别进行训练或手动标注,以识别感兴趣的区域。相比之下,SAM 的设计目标是泛化:即使从未见过特定对象和场景,它也能为任意对象和场景生成分割掩码。之前的示例。这种泛化能力源于对海量图像数据进行训练,并结合密集的标注,使得 SAM 能够根据简单的提示(例如点击、方框或文本)来隔离、勾勒或遮罩任何对象。

2025 年 11 月 19 日,Meta 发布了 Segment Anything 系列的第三代产品,名为 SAM-3:Segment Anything with Concepts(基于概念的分割)。该模型代表了分割系统在解读提示和视觉数据方面取得的重大进步。SAM-3 不再仅仅依赖几何线索或视觉示例,而是集成了可提示概念分割 (PCS)——一种允许用户使用简短的自然语言短语(例如“黄色校车”或“条纹猫”)、图像示例或它们的组合来指定语义概念的范式。作为响应,SAM-3 只需一次前向传播即可检测、分割和跟踪图像或视频中该概念的所有实例。

从下图可以看出 SAM-3 的性能,该图像被分割成“沙发”、“桌子”、“椅子”、“灯”、“地毯”和“植物”。每个分割对象都用不同的颜色显示(右图)。

SAM-3 分割示例:输入图像(左)和分割后的对象——沙发、桌子、椅子、灯、地毯、植物(右),颜色各异

3、SAM-3 分割和标注工具

以下是代码目录结构:

├── README.md
├── requirements.txt
├── sam.py                       # Wrapper that delegates to SAMSegmenter for quick runs
├── examples.py                  # Usage examples for batch and single-image processing
├── images/                      # Sample images and generated overlays (no code)
└── sam_segmentation/            # Core package
    ├── __init__.py              # Re-exports SAMSegmenter, exporters, visualizer
    ├── segmenter.py             # SAMSegmenter and SegmentationResult; main SAM-3 wrapper
    ├── exporters.py             # COCOExporter and LabelMeExporter implementations
    ├── visualizer.py            # OverlayVisualizer for rendering/saving mask overlays
    └── utils.py                 # Helpers for image I/O, EXIF handling, directory collection

抽象所有底层实现细节的核心封装器类位于 sam.py 中。

如果您想立即开始分割并了解如何使用该工具,可以直接跳到使用部分(使用封装器:简单入口点),直接查看 sam.py。

以下各节将解释此封装器的各个组成部分以及它们如何协同工作以提供一个可用于生产环境的分割流程。

我将仅解释代码的主要部分。完整的带注释的代码可在 GitHub 上找到。

4、SAMSegmenter:核心分割流程

segmenter.py 模块包含封装 SAM-3 模型加载和推理的主要分割流程。它支持多目标检测的单文本和多文本提示,管理延迟模型加载以提高内存效率,并协调导出和可视化。该模块提供单张图像和批量处理功能。

segmenter.py 中的 SAMSegmenter 类负责协调整个分割工作流程。它初始化时会配置模型路径、文本提示、导出格式和可视化设置等选项。它管理 COCO 和 LabelMe 格式的导出器,以及用于创建彩色叠加图像的可视化工具。

class SAMSegmenter:
    """Main segmentation pipeline for detecting objects using SAM-3."""

    def __init__(
        self,
        model_path: str = r"C:\Users\h02317\sam3",
        text_prompt: Union[str, List[str]] = "object",
        labels: Optional[List[str]] = None,
        export_format: Union[str, List[str]] = "coco",
        save_overlay: bool = True,
        mask_threshold: float = 0.5,
        ...
    ):
        """Initialize with configuration for model, prompts, exports, and visualization."""
        ...

    def _load_model(self):
        """Lazy load SAM-3 model on first use."""
        ...

    def _segment_image(self, image: Image.Image) -> Tuple[np.ndarray, np.ndarray, List[str]]:
        """Run SAM-3 inference on image with text prompt(s), return masks/scores/categories."""
        ...

    def process_image(self, image_path, output_dir=None) -> SegmentationResult:
        """Process single image: load → segment → export → visualize."""
        ...

    def process_directory(self, images_dir=None, output_dir=None) -> List[SegmentationResult]:
        """Batch process all images in directory, export COCO annotations."""
        ...

    def export_results(self, results, output_dir) -> Dict[str, Path]:
        """Manually export results to configured formats."""
        ...

_segment_image() 方法包含核心推理逻辑。对于单目标检测,它将文本提示转换为列表,在 SAM-3 处理器中设置图像,根据提示运行推理,并为每个生成的掩码标记类别标签。对于多目标检测,它会依次处理每个提示。对于同一幅图像,该方法会汇总所有提示的结果,并根据各自的类别标记掩码。该方法会将所有提示的掩码和分数连接成单个数组,同时保留类别关联,然后将合并后的结果以元组的形式返回。

process_directory() 方法对目录中的所有图像进行批量处理。它使用 utils.py 中的 collect_images() 实用函数收集所有图像文件,在循环开始前触发一次模型加载,然后遍历每幅图像并调用 process_image(),最后汇总所有结果。处理完所有图像后,它会将完整的数据集以 COCO 格式导出到单个 JSON 文件中。

分割器依赖于 utils.py 中的实用函数来执行常见操作。load_image_with_exif() 函数加载图像并应用 EXIF 方向校正以防止图像横向放置。collect_images() 函数扫描目录中与受支持扩展名匹配的图像文件,并返回一个排序列表以确保处理顺序的一致性。

5、结果导出:COCO 和 LabelMe 格式

推理完成并收集掩码后,需要将结果导出为标准标注格式,以便导入到标注工具或用于训练机器学习模型。exporters.py 模块提供了两个导出器类,分别用于将分割掩码转换为 COCO 和 LabelMe 格式的基于多边形的标注。

COCOExporter 类将所有分割结果导出到一个包含整个数据集的 COCO 累积 JSON 文件中。COCO 格式是一种数据集级别的格式,其中一个 JSON 文件包含所有图像、类别和标注的信息。导出器会收集所有结果中的唯一类别,创建具有顺序 ID 的类别映射,并为所有图像中检测到的每个对象生成标注条目。

class COCOExporter:
    """Export segmentation results to COCO JSON format."""

    def __init__(
        self,
        category_name: str = "object",
        dataset_name: str = "Segmentation Dataset",
        polygon_tolerance: float = 2.0,
        mask_threshold: float = 0.5
    ):
        """Initialize with dataset metadata and conversion parameters."""
        ...

    def export(
        self,
        results: List[SegmentationResult],
        output_path: Path
    ) -> dict:
        """Export list of results to single COCO JSON file."""
        ...

另一种标注格式选项是 LabelMe。LabelMeExporter 类将分割结果导出为 LabelMe 格式,该格式为每个图像创建一个 JSON 文件,而不是单个数据集文件。与使用类别 ID 和查找表的 COCO 格式不同,LabelMe 将类别标签直接存储在每个形状条目中。

class LabelMeExporter:
    """Export segmentation results to LabelMe JSON format."""

    def __init__(
        self,
        category_name: str = "object",
        polygon_tolerance: float = 2.0,
        mask_threshold: float = 0.5
    ):
        """Initialize with default category and conversion parameters."""
        ...

    def export(
        self,
        result: SegmentationResult,
        output_path: Path
    ) -> dict:
        """Export single result to LabelMe JSON file."""
        ...

两个导出器都依赖于 utils.py 中的 mask_to_polygon() 实用函数将二值分割掩码转换为多边形坐标。这种转换是必要的,因为标注工具和训练框架处理的是基于矢量的多边形,而不是基于像素的掩码。该函数使用阈值对掩码进行二值化,利用 OpenCV 查找轮廓,使用 Douglas-Peucker 算法简化多边形以减少点数并保持形状精度,并过滤掉点数少于三个的无效多边形。容差参数控制多边形精度和文件大小之间的权衡,较低的容差值会保留更多细节,但会生成更大的标注文件。

7、可视化:叠加层创建

visualizer.py 模块提供了 OverlayVisualizer 类,该类将掩码渲染为带有轮廓线的半透明彩色区域,叠加在源图像之上。

OverlayVisualizer 类负责在图像上创建彩色掩码叠加层。它维护一个包含十六种不同颜色的默认调色板,这些颜色会循环显示检测到的对象;它还支持针对特定可视化需求的自定义调色板,并提供用于程序渲染和直接保存文件的方法。

该类使用 Alpha 混合创建半透明叠加层,使原始图像和彩色蒙版同时可见,并在蒙版边界周围绘制轮廓线,以清晰地勾勒出对象边缘。

class OverlayVisualizer:
    """Create overlay visualizations of segmentation masks on images."""

    DEFAULT_COLORS = [
        (255, 0, 0), (0, 255, 0), (0, 0, 255),     # Primary RGB colors
        (255, 255, 0), (255, 0, 255), (0, 255, 255), # Secondary colors
        ...  # 16 distinct colors total
    ]

    def __init__(
        self,
        alpha: float = 0.4,
        colors: Optional[List[Tuple[int, int, int]]] = None,
        contour_thickness: int = 2,
        draw_labels: bool = False,
        mask_threshold: float = 0.5
    ):
        """Initialize with visualization parameters for colors, transparency, and styling."""
        ...

    def render(
        self,
        image: Union[np.ndarray, Image.Image],
        masks: np.ndarray,
        labels: Optional[List[str]] = None
    ) -> np.ndarray:
        """Core rendering: blend colored masks onto image with contours and optional labels."""
        ...

    def save(
        self,
        result: SegmentationResult,
        output_path: Path,
        labels: Optional[List[str]] = None
    ):
        """Generate and save overlay from SegmentationResult object."""
        ...

    def save_from_path(
        self,
        image_path: Path,
        masks: np.ndarray,
        output_path: Path,
        labels: Optional[List[str]] = None
    ):
        """Generate and save overlay directly from image path and masks."""
        ...

render() 方法包含核心可视化逻辑。它首先将输入图像转换为 NumPy 数组(如果它是 PIL 图像),并确保它是 RGB 格式。对于输入数组中的每个掩码,它使用阈值对掩码进行二值化,通过模运算循环遍历调色板选择颜色,并应用 Alpha 混合将掩码颜色与原始图像像素混合。

8、使用封装器:简易入口

sam.py 文件提供了一个简单的入口,演示如何使用 SAMSegmenter 类进行单张或批量图像分割。要在您自己的代码中使用此封装器,只需导入该类并使用所需的参数进行配置即可。您只需以下代码片段即可运行整个工作流程。

from sam_segmentation import SAMSegmenter

# Initialize segmenter with configuration
segmenter = SAMSegmenter(
    text_prompt=["Segment each tooth accurately"],  # Prompt for SAM-3 inference
    labels=["tooth"],                                # Label for annotations
    export_format=["coco"],                            # Export format(s)
    save_overlay=True,                                 # Generate visualizations
    images_dir="./images",                             # Input directory
    output_dir="./images",                             # Output directory
)

# Process all images in the directory
results = segmenter.process_directory()

其他参数控制阈值、可视化外观、多边形精度以及分割流程的其他方面。有关所有可用参数及其说明的完整列表,请参阅 README 文件,其中提供了配置选项的文档。

另请参阅 examples.py,了解如何使用不同的配置。

9、结果

在构建了用于处理分割和标注的 SAM-3 Python 封装程序后,我运行的第一个测试是牙齿分割,即使是零样本测试的结果也令人印象深刻。该模型生成了精确的掩膜,我可以将其导出为 COCO JSON 文件并直接加载到标注工具中。

输入图像(左)和带有分割牙齿的叠加图像(右)
输入图像(左)和带有分割牙齿的叠加图像(右)

以下是其中一张分割图像在 cvat.ai 标注工具中的显示效果。边界与实际牙齿轮廓对齐,几乎不需要手动校正。测试集中的大多数分割结果已经非常好了,这意味着我可以快速检查任何需要微调的分割结果,并在工具中进行优化。

导入到 CVAT标注工具中的图像和预标注数据 (COCO)

此工作流程可让您快速构建数据集。您无需从头开始绘制每个多边形,而是从 SAM-3 的输出开始,仅在必要时进行修正。这使得准备用于训练新模型的数据变得容易,甚至可以微调 SAM-3 本身以获得更好的性能。

我还用其他图像进行了测试。以下是我使用的提示信息的结果。

text_prompt=["road", "sidewalk", "building"]
输入图像(左)和分割出道路、人行道和建筑物的叠加图像(右)
text_prompt=["Segment each bus accurately", "Segment each delivery truck accurately"]
输入图像(左)和分割出所有公交车和送货卡车的叠加图像(右)
text_prompt=["pedestrian"]

10、进一步探索

SAM-3 可以用作 AI 代理中的工具,用于评估和改进自身的输出。与一次性生成分割结果后就停止不同,智能体工作流程可以反复测试分割结果与提示的匹配度,评估质量指标,并调整参数(例如掩码阈值或提示措辞),直到达到所需的质量。这对于提示信息较为复杂的情况非常有用。plex。

我发现分割质量也取决于提示语。有些提示语能生成精确完整的掩码,而有些则会产生零散或不一致的输出。智能体可以利用这一点,将提示语动态拆分成更具体的子提示语,尝试同义词,或重新表述模糊的概念,以找到最佳结果。


原文链接:How I Built a Tool to "Literally" Segment Anything in Images and Automatically Prepare a Labelled…

汇智网翻译整理,转载请标明出处