开源项目实战学习之YOLO11：ultralytics-cfg-models-fastsam（九）

👉 点击关注不迷路

文章大纲

- [1. init.py](#1. init.py)
- [2. model.py](#2. model.py)
- [3. predict.py](#3. predict.py)
- [4. utils.py](#4. utils.py)
- [5. val.py](#5. val.py)

FastSAM 是一种目标检测和图像分割模型，Ultralytics 是一个在计算机视觉领域广泛使用的库，用于各种深度学习模型的训练、推理和评估等任务。
ultralytics-cfg-models-fastsam 这个路径下可能包含了 FastSAM 模型的配置文件，这些配置文件用于定义模型的结构、超参数（如学习率、批次大小、训练轮数等）以及数据预处理和后处理的方式等。
同时，该路径也可能包含已经训练好的 FastSAM 模型权重文件，以便在进行推理或进一步微调时使用。
通过这些配置和模型文件，用户可以方便地使用 Ultralytics 库来加载 FastSAM 模型，进行图像分割任务，例如对输入的图像进行目标检测和分割，识别出图像中的不同物体并为其生成相应的分割掩码。
此外，也可以基于这些配置和模型文件进行模型的训练和优化，以适应不同的数据集和应用场景。

1. init.py

python 复制代码

from .model import FastSAM  
# 从当前目录的model模块中导入FastSAM类

from .predict import FastSAMPredictor  
# 从当前目录的predict模块中导入FastSAMPredictor类

from .val import FastSAMValidator  
# 从当前目录的val模块中导入FastSAMValidator类

# 定义模块的公共接口，即可以通过from module import *导入的名称列表
__all__ = "FastSAMPredictor", "FastSAM", "FastSAMValidator"

2. model.py

FastSAM 即 Fast Segment Anything Model

一种高效的图像分割模型，依托于 Ultralytics 的开发生态，在多领域展现出独特价值
- 模型架构与原理： 它可能基于 Transformer 架构，通过对图像特征的深度挖掘，实现对各类物体的精准分割。在处理复杂场景图像时，能有效捕捉物体的边缘和细节信息，以较低的计算成本快速生成高质量的分割结果。
功能特点
- 快速处理： 相比传统的图像分割模型，FastSAM 在保证分割精度的同时，大幅提升了处理速度。在实时性要求较高的场景，如自动驾驶的道路场景分割、直播内容的实时物体分割等，能快速处理图像或视频流，满足实时性需求。
- 多模态提示支持： 支持多种提示方式进行分割，如边界框（bounding boxes）、点（points）、标签（labels）和文本（texts）。用户可根据具体需求，灵活选择提示信息，引导模型对特定目标进行分割。
- 通用性强： 可应用于多种图像分割任务，包括但不限于实例分割、语义分割和全景分割。无论是自然场景图像、医学影像，还是工业检测图像，都能展现出良好的分割性能。
应用场景
- 计算机视觉研究： 为研究人员提供了一个高效的图像分割工具，可用于探索新的分割算法、验证研究思路。在新型神经网络架构的研究中，利用 FastSAM 快速获取分割结果，评估架构的有效性。
- 自动驾驶： 用于识别道路上的车辆、行人、交通标志等目标，为自动驾驶汽车的决策提供关键信息。通过实时分割道路场景图像，帮助车辆准确感知周围环境，实现安全行驶。
- 医学影像分析： 在医学领域，能辅助医生对医学影像（如 X 光、CT、MRI 等）进行分析。帮助医生快速分割出病变组织、器官等感兴趣区域，提高诊断效率和准确性。
- 工业检测： 在工业生产中，对产品表面缺陷进行检测时，可分割出缺陷区域，判断产品是否合格。对电子芯片、机械零部件等进行质量检测，及时发现生产过程中的问题。

python 复制代码

from pathlib import Path
# 导入 Path 类，用于处理文件路径

from ultralytics.engine.model import Model
# 从 ultralytics 引擎模块导入 Model 基类

from .predict import FastSAMPredictor
# 从当前包中导入 FastSAMPredictor 类，用于进行预测操作

from .val import FastSAMValidator
# 从当前包中导入 FastSAMValidator 类，用于进行验证操作

class FastSAM(Model):
    """
    FastSAM model interface for segment anything tasks.
    # FastSAM 模型接口，用于处理任意图像分割任务

    # 该类继承自 Model 基类，为 FastSAM（快速任意分割模型）实现提供特定功能，可实现高效且准确的图像分割

    Attributes:
        model (str): Path to the pre - trained FastSAM model file.
        # 预训练的 FastSAM 模型文件的路径
        task (str): The task type, set to "segment" for FastSAM models.
        # 任务类型，对于 FastSAM 模型，设置为 "segment"（分割）

    Examples:
        >>> from ultralytics import FastSAM
        >>> model = FastSAM("last.pt")
        >>> results = model.predict("ultralytics/assets/bus.jpg")
    # 使用示例，展示如何导入 FastSAM 类、初始化模型并进行预测
    """

    def __init__(self, model="FastSAM-x.pt"):
        """
        Initialize the FastSAM model with the specified pre - trained weights.
        # 使用指定的预训练权重初始化 FastSAM 模型

        Args:
            model (str): Path to the pre - trained FastSAM model file. Defaults to "FastSAM-x.pt".
            # 预训练的 FastSAM 模型文件的路径，默认为 "FastSAM-x.pt"
        """
        if str(model) == "FastSAM.pt":
            model = "FastSAM-x.pt"
            # 如果传入的模型名称是 "FastSAM.pt"，则将其替换为 "FastSAM-x.pt"
        assert Path(model).suffix not in {".yaml", ".yml"}, "FastSAM models only support pre - trained models."
        # 断言传入的模型文件后缀不是 .yaml 或 .yml，因为 FastSAM 模型仅支持预训练模型
        
        super().__init__(model=model, task="segment")
        # 调用父类 Model 的构造函数，传入模型路径和任务类型

    def predict(self, source, stream=False, bboxes=None, points=None, labels=None, texts=None, **kwargs):
        """
        # 对图像或视频源进行分割预测

        # 支持使用边界框、点、标签和文本进行提示分割。该方法将这些提示信息打包并传递给父类的 predict 方法

        Args:
            source (str | PIL.Image | numpy.ndarray): Input source for prediction, can be a file path, URL, PIL image,
                or numpy array.
            # 预测的输入源，可以是文件路径、URL、PIL 图像或 numpy 数组
            stream (bool): Whether to enable real - time streaming mode for video inputs.
            # 是否为视频输入启用实时流模式
            # 用于提示分割的边界框坐标，格式为 [[x1, y1, x2, y2], ...]
            # 用于提示分割的点坐标，格式为 [[x, y], ...]
            # 用于提示分割的类别标签
            # 用于分割引导的文本提示
            # 传递给预测器的其他关键字参数

        Returns:
            # 包含预测结果的 Results 对象列表
        """
        
        prompts = dict(bboxes=bboxes, points=points, labels=labels, texts=texts)
        # 将边界框、点、标签和文本提示信息打包成字典
        
        return super().predict(source, stream, prompts=prompts, **kwargs)
        # 调用父类的 predict 方法，传入输入源、流模式、提示信息和其他关键字参数，并返回预测结果

    @property
    def task_map(self):
        """
        # 返回一个字典，将分割任务映射到相应的预测器和验证器类
        """
        return {"segment": {"predictor": FastSAMPredictor, "validator": FastSAMValidator}}
        # 返回一个字典，键为 "segment"，值为包含预测器和验证器类的字典

3. predict.py

关键词：图像分割预测、边界框的交并比、缩放掩码
CLIP模型
- CLIP（Contrastive Language-Image Pretraining）模型是 OpenAI 开发的一种开创性的神经网络，通过互联网上大量多样的（图像，文本）对进行训练，具备强大的跨模态理解能力，能够将自然语言与图像信息紧密联系起来。
- 模型架构
  - 图像编码器： 可选用 Vision Transformer（ViT）或 ResNet 等架构。以 ViT 为例，它将图像划分为多个小块，然后像处理文本序列一样处理这些图像块，通过多头注意力机制捕捉图像的全局特征 。
  - 文本编码器： 基于文本 Transformer，把文本转换为连续的向量表示，在这个过程中理解文本语义和结构信息。
  - 共享嵌入空间： 两个编码器将图像和文本投影到同一个向量空间，在这个空间中，语义相似的图像和文本对其向量距离更近，为后续的匹配任务奠定基础。
- 零样本学习能力： CLIP 最显著的优势是零样本学习。在 ImageNet 分类任务中，它无需使用 ImageNet 训练集中 128 万张标记示例进行训练，就能达到与原始 ResNet50 模型相匹配的性能。使用时，只要提供文本描述（如 "一只猫""一辆汽车"），CLIP 模型就能对图像进行分类，判断图像内容是否与文本匹配。

ViT-B/32

ViT-B/32 是 CLIP 模型中使用的一种视觉 Transformer（Vision Transformer，ViT）架构的具体变体。
模型结构：
- "ViT"： 代表视觉 Transformer，是一种将 Transformer 架构应用于计算机视觉任务的模型。它将图像分割成一系列的图像块（patches），并将这些图像块作为输入序列，类似于自然语言处理中 Transformer 对文本序列的处理方式。通过这种方式，ViT 可以有效地学习图像中的全局信息和长期依赖关系。
- "B"： 通常表示基础（Base）版本，指的是模型的规模和复杂度处于中等水平。例如，在参数数量、层数、隐藏层维度等方面，基础版本具有一定的设定，是一种相对较为平衡的模型配置，既能够在性能和计算资源之间取得较好的权衡，又能在多种视觉任务上取得不错的效果。
- "32"： 表示图像块的大小为 32×32 像素。这意味着在将图像输入到 ViT 模型之前，会先将图像分割成边长为 32 像素的正方形小块。较小的图像块大小可以捕捉到更精细的图像细节，但也会增加模型的计算量和参数量；而较大的图像块大小则可以减少计算量，但可能会丢失一些细节信息。

python 复制代码

	import torch
	# 导入PyTorch库，用于深度学习相关的张量计算和模型操作
	
	from PIL import Image
	# 导入PIL库的Image模块，用于处理图像
	
	from ultralytics.models.yolo.segment import SegmentationPredictor
	# 从ultralytics的yolo模型的segment模块中导入SegmentationPredictor类，可能是用于图像分割预测的基类
	
	from ultralytics.utils import DEFAULT_CFG, checks
	# 从ultralytics的utils模块中导入DEFAULT_CFG（可能是默认配置）和checks（可能用于检查某些条件或配置）
	
	from ultralytics.utils.metrics import box_iou
	# 从ultralytics的utils模块的metrics子模块中导入box_iou函数，可能用于计算边界框的交并比
	
	from ultralytics.utils.ops import scale_masks
	# 从ultralytics的utils模块的ops子模块中导入scale_masks函数，可能用于缩放掩码
	
	from .utils import adjust_bboxes_to_image_border
	# 从当前目录的utils模块中导入adjust_bboxes_to_image_border函数，可能用于调整边界框以适应图像边界
	
	class FastSAMPredictor(SegmentationPredictor):
		# FastSAMPredictor类，继承自SegmentationPredictor，用于图像分割预测，并支持多种提示方式。
	    def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
	        """ 初始化函数。
	        Args:
	            cfg (dict, optional): 配置字典，默认为DEFAULT_CFG。
	            overrides (dict, optional): 用于覆盖默认配置的字典。
	            _callbacks (list, optional): 回调函数列表。
	        """
	        super().__init__(cfg, overrides, _callbacks)
	        # 初始化一个空字典，用于存储各种提示信息（边界框、点、标签、文本等）
	        self.prompts = {}
	
	    def postprocess(self, preds, img, orig_imgs):
	        """ 对模型预测结果进行后处理。
	        Args:
	            preds (torch.Tensor): 模型的预测结果。
	            img (torch.Tensor): 输入的图像张量。
	            orig_imgs (list): 原始图像列表。
	        Returns:
	            list: 经过后处理和提示应用后的结果列表。
	        """
	        # 从prompts字典中弹出【边界框】提示信息，如果不存在则返回None
	        bboxes = self.prompts.pop("bboxes", None)
	        
	        # 从prompts字典中弹出【点】提示信息，如果不存在则返回None
	        points = self.prompts.pop("points", None)
	        
	        # 从prompts字典中弹出【标签】提示信息，如果不存在则返回None
	        labels = self.prompts.pop("labels", None)
	        
	        # 从prompts字典中弹出【文本】提示信息，如果不存在则返回None
	        texts = self.prompts.pop("texts", None)
	
	        # 调用父类的postprocess方法进行基本的后处理
	        results = super().postprocess(preds, img, orig_imgs)
	
	        for result in results:
	            # 创建一个表示整个图像边界的张量，格式为 [x1, y1, x2, y2]
	            full_box = torch.tensor(
	                [0, 0, result.orig_shape[1], result.orig_shape[0]], device=preds[0].device, dtype=torch.float32
	            )
	            
	            # 调整预测的边界框，使其适应原始图像的边界
	            boxes = adjust_bboxes_to_image_border(result.boxes.xyxy, result.orig_shape)
	            
	            # 计算【全图像边界框与调整后的边界框】之间的交并比，找到【交并比大于0.9的索引】
	            idx = torch.nonzero(box_iou(full_box[None], boxes) > 0.9).flatten()
	            if idx.numel() != 0:
	                # 如果存在【交并比大于0.9的边界框，则将其设置为全图像边界】
	                result.boxes.xyxy[idx] = full_box
	
	        # 调用prompt方法，应用各种提示信息到结果中
	        return self.prompt(results, bboxes=bboxes, points=points, labels=labels, texts=texts)
	
	    def prompt(self, results, bboxes=None, points=None, labels=None, texts=None):
	        """
	        根据提供的提示信息（边界框、点、标签、文本）对分割结果进行筛选和处理。
	        Args:
	            results (list or object): 分割结果，可以是【单个结果或结果列表】。
	            bboxes (list, optional): 边界框提示信息，格式为 [[x1, y1, x2, y2], ...]。
	            points (list, optional): 点提示信息，格式为 [[x, y], ...]。
	            labels (list, optional): 标签提示信息，与点提示信息对应。
	            texts (list, optional): 文本提示信息。
	        Returns:
	            list: 经过提示筛选后的结果列表。
	        """
	        if bboxes is None and points is None and texts is None:
	            # 如果没有提供任何提示信息，则直接返回原始结果
	            return results
	
	        prompt_results = []
	        if not isinstance(results, list):
	            # 如果结果不是列表，则将其转换为列表
	            results = [results]
	
	        for result in results:
	            if len(result) == 0:
	                # 如果结果为空，则直接添加到提示结果列表中
	                prompt_results.append(result)
	                continue
	
	            masks = result.masks.data
	            if masks.shape[1:] != result.orig_shape:
	                # 如果【掩码的形状与原始图像形状不一致】，则缩放掩码
	                masks = scale_masks(masks[None], result.orig_shape)[0]
	
	            # 初始化一个布尔张量，用于标记符合条件的分割结果
	            idx = torch.zeros(len(result), dtype=torch.bool, device=self.device)
	
	            if bboxes is not None:
	                # 将边界框提示信息转换为张量
	                bboxes = torch.as_tensor(bboxes, dtype=torch.int32, device=self.device)
	                if bboxes.ndim == 1:
	                    # 如果边界框是一维的，则将其转换为二维
	                    bboxes = bboxes[None]
	
	                # 计算每个边界框的面积
	                bbox_areas = (bboxes[:, 3] - bboxes[:, 1]) * (bboxes[:, 2] - bboxes[:, 0])
	                # 计算每个掩码与边界框重叠部分的面积
	                mask_areas = torch.stack([masks[:, b[1] : b[3], b[0] : b[2]].sum(dim=(1, 2)) for b in bboxes])
	                
	                # 计算每个掩码的总面积
	                full_mask_areas = torch.sum(masks, dim=(1, 2))
	
	                # 计算边界框与掩码的并集面积
	                union = bbox_areas[:, None] + full_mask_areas - mask_areas
	                
	                # 找到重叠面积与并集面积比值最大的索引，并将其对应的idx位置设为True
	                idx[torch.argmax(mask_areas / union, dim=1)] = True
	
	            if points is not None:
	                # 将点提示信息转换为张量
	                points = torch.as_tensor(points, dtype=torch.int32, device=self.device)
	                if points.ndim == 1:
	                    # 如果点是一维的，则将其转换为二维
	                    points = points[None]
	
	                if labels is None:
	                    # 如果没有提供标签，则创建全为1的标签张量
	                    labels = torch.ones(points.shape[0])
	
	                # 将标签转换为张量
	                labels = torch.as_tensor(labels, dtype=torch.int32, device=self.device)
	                assert len(labels) == len(points), (
	                    f"Expected `labels` to have the same size as `point`, but got {len(labels)} and {len(points)}"
	                )
	
	                # 根据标签的总和初始化point_idx，如果【标签总和为0（即全为负点）】，则设为全True，否则设为全False
	                point_idx = (
	                    torch.ones(len(result), dtype=torch.bool, device=self.device)
	                    if labels.sum() == 0
	                    else torch.zeros(len(result), dtype=torch.bool, device=self.device)
	                )
	
	                for point, label in zip(points, labels):
	                    # 根据点的位置和标签，更新point_idx
	                    point_idx[torch.nonzero(masks[:, point[1], point[0]], as_tuple=True)[0]] = bool(label)
	                    
	                # 将point_idx与idx进行逻辑或操作
	                idx |= point_idx
	
	            if texts is not None:
	                if isinstance(texts, str):
	                    # 如果文本提示是字符串，则将其转换为列表 ？？？
	                    texts = [texts]
	
	                crop_ims, filter_idx = [], []
	                for i, b in enumerate(result.boxes.xyxy.tolist()):
	                    x1, y1, x2, y2 = (int(x) for x in b)
	                    if masks[i].sum() <= 100:
	                        # 如果掩码的总和小于等于100，则将其索引添加到filter_idx中并跳过
	                        filter_idx.append(i)
	                        continue
	                    # 从原始图像中裁剪出边界框对应的区域，并转换为PIL图像
	                    crop_ims.append(Image.fromarray(result.orig_img[y1:y2, x1:x2, ::-1]))
	                # 使用CLIP模型计算裁剪图像与文本提示之间的相似度
	                similarity = self._clip_inference(crop_ims, texts)
	                
	                # 找到相似度最大的索引
	                text_idx = torch.argmax(similarity, dim=-1)
	
	                if len(filter_idx):
	                    # 如果存在过滤索引，则调整text_idx
	                    text_idx += (torch.tensor(filter_idx, device=self.device)[None] <= int(text_idx)).sum(0)
	                
	                # 将text_idx对应的idx位置设为True
	                idx[text_idx] = True
	
	            # 将符合条件的结果添加到提示结果列表中
	            prompt_results.append(result[idx])
	
	        return prompt_results
	
	    def _clip_inference(self, images, texts):
	        """ 使用CLIP模型进行推理，计算图像与文本之间的相似度。
	        Args:
	            images (list): PIL图像列表。
	            texts (list): 文本提示列表。
	        Returns: torch.Tensor: 图像与文本之间的相似度矩阵，形状为 (M, N)，M为文本数量，N为图像数量。
	        """
	        try:
	            import clip
	        except ImportError:
	            # 如果CLIP库未安装，则检查并安装
	            checks.check_requirements("git+https://github.com/ultralytics/CLIP.git")
	            import clip
	
	        if (not hasattr(self, "clip_model")) or (not hasattr(self, "clip_preprocess")):
	            # 如果当前对象没有CLIP模型和预处理函数，则加载CLIP模型和预处理函数
	            self.clip_model, self.clip_preprocess = clip.load("ViT-B/32", device=self.device)
	
	        # 对图像进行预处理，并将其转换为张量
	        images = torch.stack([self.clip_preprocess(image).to(self.device) for image in images])
	        
	        # 对文本进行分词，并将其转换为张量
	        tokenized_text = clip.tokenize(texts).to(self.device)
	
	        # 使用CLIP模型对图像进行编码，得到图像特征
	        image_features = self.clip_model.encode_image(images)
	        
	        # 使用CLIP模型对文本进行编码，得到文本特征
	        text_features = self.clip_model.encode_text(tokenized_text)
	
	        # 对图像特征进行归一化
	        image_features /= image_features.norm(dim=-1, keepdim=True)
	        
	        # 对文本特征进行归一化
	        text_features /= text_features.norm(dim=-1, keepdim=True)
	
	        # 计算图像特征与文本特征之间的相似度，并返回相似度矩阵
	        return (image_features * text_features[:, None]).sum(-1)
	
	    def set_prompts(self, prompts):
	        """
	        设置提示信息字典。
	        Args:
	            prompts (dict): 包含各种提示信息的字典，如 "bboxes", "points", "labels", "texts" 等。
	        """
	        self.prompts = prompts
	```

4. utils.py

python 复制代码

def adjust_bboxes_to_image_border(boxes, image_shape, threshold=20):
    """
    将边界框调整到图像边界附近，确保边界框不会超出合理范围。
    Args:
        boxes (torch.Tensor或numpy.ndarray): 
        边界框的张量或数组，形状通常为 (N, 4)，N表示边界框的数量，
        每个边界框包含四个坐标值 (x1, y1, x2, y2)，分别代表左上角和右下角的坐标。
        image_shape (tuple): 图像的形状，格式为 (高度, 宽度)。
        threshold (int, 可选): 接近边界的阈值。
        	如果边界框的坐标值与图像边界的距离小于该阈值，
        	则将边界框的坐标调整到边界上。默认为20。
    Returns:
        torch.Tensor或numpy.ndarray: 调整后的边界框张量或数组，形状与输入的boxes相同。
    """
    # 获取图像的高度和宽度
    h, w = image_shape

    # 调整靠近图像左边界的边界框的x1坐标
    boxes[boxes[:, 0] < threshold, 0] = 0  # x1
    
    # 调整靠近图像上边界的边界框的y1坐标
    boxes[boxes[:, 1] < threshold, 1] = 0  # y1
    
    # 调整靠近图像右边界的边界框的x2坐标
    boxes[boxes[:, 2] > w - threshold, 2] = w  # x2
    
    # 调整靠近图像下边界的边界框的y2坐标
    boxes[boxes[:, 3] > h - threshold, 3] = h  # y2
    return boxes

5. val.py

python 复制代码

from ultralytics.models.yolo.segment import SegmentationValidator  
# 从ultralytics的yolo模型的segment模块导入SegmentationValidator类

from ultralytics.utils.metrics import SegmentMetrics  
# 从ultralytics的utils模块的metrics子模块导入SegmentMetrics类


class FastSAMValidator(SegmentationValidator):  
# 定义FastSAMValidator类，继承自SegmentationValidator

    def __init__(self, dataloader=None, save_dir=None, pbar=None, args=None, _callbacks=None):
        """
        初始化FastSAMValidator类的实例。
        Args:
            dataloader (DataLoader, 可选): 数据加载器，用于加载验证数据。默认为None。
            
            save_dir (str, 可选): 保存验证结果的目录。默认为None。
            
            pbar (tqdm.tqdm, 可选): 进度条对象，用于显示验证过程的进度。默认为None。
            
            args (Namespace, 可选): 包含验证参数的命名空间。默认为None。
            
            _callbacks (list, 可选): 回调函数列表，用于在验证过程中执行特定的操作。默认为None。
        """
        super().__init__(dataloader, save_dir, pbar, args, _callbacks)  # 调用父类的初始化方法
        self.args.task = "segment"  # 设置任务类型为"segment"（分割）
        
        self.args.plots = False  # 禁用混淆矩阵和其他绘图，以避免错误
        
        self.metrics = SegmentMetrics(save_dir=self.save_dir)  
        # 创建SegmentMetrics对象，用于计算分割任务的指标，并指定保存目录