FastSAM是一个在计算机视觉领域获得显著关注的开源图像分割模型。它以其准确性迅速获得了社区的认可。然而,SAM模型在实际应用中受限于其对计算资源的高需求,特别是在需要实时处理的场景中。FastSAM通过采用一种解耦的方法,有效地克服了这一挑战,它将图像分割任务分为两个连续的阶段:全实例分割和提示引导的选择。
FastSAM通过使用一个分离的方法来降低计算需求,它将分割任务分解为两个阶段:全实例分割和提示引导的选择。这种方法使得FastSAM在运行速度上比SAM快50倍,同时保持了相当的准确性。FastSAM的训练仅使用了SAM训练数据集的2%,这一技术被称为数据集蒸馏,使得FastSAM在运行速度上大幅提升,同时在性能上与SAM相当。
本指南将展示如何结合使用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是SAM的一个强大补充,用于图像分割任务。FastSAM通过采用一种解耦的方法,并利用基于卷积神经网络(CNN)的检测器,克服了SAM的计算限制。这使得FastSAM能够在不显著牺牲性能质量的情况下实现实时分割。
通过仅在SA-1B数据集的2%上训练CNN检测器,FastSAM实现了与SAM相当的性能,同时运行速度提高了50倍。这使得FastSAM更适合于速度至关重要的场景。