FastSAM图像分割模型介绍

FastSAM是一个在计算机视觉领域获得显著关注的开源图像分割模型。它以其准确性迅速获得了社区的认可。然而,SAM模型在实际应用中受限于其对计算资源的高需求,特别是在需要实时处理的场景中。FastSAM通过采用一种解耦的方法,有效地克服了这一挑战,它将图像分割任务分为两个连续的阶段:全实例分割和提示引导的选择。

FastSAM的优势

FastSAM通过使用一个分离的方法来降低计算需求,它将分割任务分解为两个阶段:全实例分割和提示引导的选择。这种方法使得FastSAM在运行速度上比SAM快50倍,同时保持了相当的准确性。FastSAM的训练仅使用了SAM训练数据集的2%,这一技术被称为数据集蒸馏,使得FastSAM在运行速度上大幅提升,同时在性能上与SAM相当。

如何使用FastSAM

本指南将展示如何结合使用SAM和FastSAM来直观比较它们的性能。如果只对FastSAM感兴趣,可以跳过安装和使用SAM的部分。以下是安装FastSAM和SAM及其依赖项的步骤:

!git clone https://github.com/CASIA-IVA-Lab/FastSAM.git !pip -q install -r FastSAM/requirements.txt !pip -q install git+https://github.com/openai/CLIP.git roboflow supervision !pip -q install git+https://github.com/facebookresearch/segment-anything.git !wget -P FastSAM/weights https://huggingface.co/spaces/An-619/FastSAM/resolve/main/weights/FastSAM.pt !wget -P FastSAM/weights https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth

接下来,将导入所需的库并加载FastSAM模型

from fastsam import FastSAM, FastSAMPrompt from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor import supervision as sv import roboflow from roboflow import Roboflow model = FastSAM('./weights/FastSAM.pt')

现在,将使用FastSAM在一些示例图像上可视化生成的分割掩码。代码段将遍历示例图像列表,并使用FastSAM生成分割掩码。然后,结果掩码将被可视化并保存为输出图像。

retina_masks=True imgsz=1024 conf=0.4 iou=0.9 folder = './images/' Images = ['tool1.jpg', 'bone.jpg', 'stamp.jpg', 'plant1.jpg', 'chip.jpeg', 'pill1.png'] for idx, img_name in enumerate(Images): path = os.path.join(folder, img_name) everything_results = model(path, device=DEVICE, retina_masks=True, imgsz=1024, conf=0.4, iou=0.9) prompt_process = FastSAMPrompt(path, everything_results, device=DEVICE) ann = prompt_process.everything_prompt() output_filename = f'output_{idx}.jpg' output_path = os.path.join('./output/', output_filename) prompt_process.plot(annotations=ann, output=output_path)

为了进一步展示FastSAM的能力,将把它应用到Roboflow 100基准数据集上。将使用文本提示来引导分割过程。首先,需要从Roboflow Universe下载Roboflow基准数据集,特别是训练集中的图像。

roboflow.login() rf = Roboflow() project = rf.workspace("roboticfish").project("underwater_object_detection") dataset = project.version(8).download("yolov8") train_folder = os.path.join(dataset.location, 'train', 'images')

现在,将使用文本提示将FastSAM应用于Roboflow基准数据集中的一张图像。在这个例子中,将选择一张图像并提供提示“Penguin”来引导分割过程。

everything_results = model(image_path, device=DEVICE, retina_masks=True, imgsz=1024, conf=0.4, iou=0.9,) prompt_process = FastSAMPrompt(image_path, everything_results, device=DEVICE) ann = prompt_process.text_prompt(text='Penguin') prompt_process.plot(annotations=ann, output='./output/')

为了比较SAM和FastSAM生成的分割掩码,将提取两个模型的掩码并进行可视化。将使用everything_prompt()函数来提取FastSAM的分割掩码,并将结果掩码可视化并保存为输出图像。

IMAGE_PATH = './images/real3.jpeg' everything_results = model(IMAGE_PATH, device=DEVICE, retina_masks=True, imgsz=1024, conf=0.4, iou=0.9,) prompt_process = FastSAMPrompt(IMAGE_PATH, everything_results, device=DEVICE) ann = prompt_process.everything_prompt() prompt_process.plot(annotations=ann, output='./output/')

在下面的代码片段中,使用SAM为相同的示例图像提取分割掩码。结果掩码存储在sam_masks变量中。

image = cv2.cvtColor(cv2.imread(IMAGE_PATH), cv2.COLOR_BGR2RGB) sam_checkpoint = "sam_vit_h_4b8939.pth" model_type = "vit_h" sam = sam_model_registry[model_type](checkpoint=sam_checkpoint) sam.to(device=DEVICE) mask_generator = SamAutomaticMaskGenerator(sam) sam_masks = mask_generator.generate(image)

下面的代码比较了FastSAM和SAM生成的分割掩码。FastSAM的输出图像和带有SAM掩码的注释图像并排显示,以进行视觉比较。

fastsam_output = cv2.imread("./output/real3.jpeg") image_bgr = cv2.imread(IMAGE_PATH) mask_annotator = sv.MaskAnnotator() detections = sv.Detections.from_sam(sam_result=sam_masks) annotated_image = mask_annotator.annotate(scene=image.copy(), detections=detections) sv.plot_images_grid( images=[fastsam_output, annotated_image], grid_size=(1, 2), titles=['FastSAM segmented image', 'SAM segmented image'] )

通过比较FastSAM和SAM生成的掩码来可视化不同的掩码。当交集比率(IOU)低于0.05时,认为掩码是不同的,并将其显示在原始图像上方。

FastSAM的推理结果

FastSAM是SAM的一个强大补充,用于图像分割任务。FastSAM通过采用一种解耦的方法,并利用基于卷积神经网络(CNN)的检测器,克服了SAM的计算限制。这使得FastSAM能够在不显著牺牲性能质量的情况下实现实时分割。

通过仅在SA-1B数据集的2%上训练CNN检测器,FastSAM实现了与SAM相当的性能,同时运行速度提高了50倍。这使得FastSAM更适合于速度至关重要的场景。

沪ICP备2024098111号-1
上海秋旦网络科技中心:上海市奉贤区金大公路8218号1幢 联系电话:17898875485