CV | Segment Anything论文详解及代码实现

本文主要是详解解释了SAM的论文部分以及代码实现~

论文:2023.04.05_Segment Anything

论文地址:2304.02643.pdf (arxiv.org)

代码地址:facebookresearch/segment-anything: The repository provides code for running inference with the SegmentAnything Model (SAM), links for downloading the trained model checkpoints, and example notebooks that show how to use the model. (github.com)

1.论文详解

SAM 包括三个主要组件:图像编码器、提示编码器和掩蔽解码器。

图像编码器由基于 MAE 预训练的 ViT 构成,而掩蔽解码器如图所示

图像编码器(Image encoder)

基于ViT模型对于图像编码作为输入,利用MAE(Masked AutoEncoder)预训练的ViT模型,对每张图片只处理一次。输入(c,h,w)的图像,对图像进行缩放,按照长边缩放成1024,短边不够就填充(pad),得到(c,1024,1024)的图像,经过image encoder,得到对图像16倍下采样的feature,大小为(256,64,64)。

提示编码器(Prompt encoder )

俩种提示:

稀疏sparse

  • 点points:使用position encodings。点映射到256维的向量,代表点位置的 positional encoding,加2个代表该点是前景/背景的可学习的embedding
  • boxes:使用position encodings,俩点确定一个box,用一个embedding对表示:1) 可学习的embedding代表左上角,2) 可学习的embedding代表右下角
  • text:使用CLIP作为encoder

密集dense(mask):使用卷积作为encoder,mask 和image embedding通过element-wise相乘 (逐元素相乘,可以理解成mask的feature对image的feature进行加权)。在内部会缩小4倍,通过2个2×2,步幅为2的卷积实现,输出的通道为4和16.然后1×1卷积将这些通道映射到256个通道,每一次都通过层归一化进行增强,然后将蒙版按元素添加到图像嵌入中。在未提示掩码提示的情况下,将每个图像嵌入位置添加一个表示无掩码的学习嵌入。

掩码解码(Mask decoder)

掩码解码器有效地将图像嵌入、提示嵌入和输出令牌映射到掩码。它使用改进的 Transformer 解码器,后跟动态掩模预测头。输出token作为掩码,运行俩个块后,对图像编码和MLPt图输出到动态线性分类(dynamic linear classifier),用来计算掩码区域在每张图像的位置.

 每个解码器层执行 4 个步骤(从上图可以看出):

1)对 output tokens + prompt tokens 进行 self-attention;
2)将 1)得到的 tokens 作为 query token 与 image embedding 做 cross-attention;
3)使用 MLP 更新 2)得到的 tokens;
4)将 image embedding 作为 query token 与 3)得到的 tokens 做 cross-attention。

 最后一步使用提示信息更新图像嵌入。在交叉注意力期间,图像嵌入被视为一组 64×64×256 维向量。

推理时,sam模型是被封装在SamPredictor类中,将sam的forward预测的流程分别拆解到SamPredictor类的不同方法中、分不同阶段进行。
sam中forward函数对Image encoder、Prompt encoder和Mask decoder三个操作是连续的。

2.项目实现

docker 容器(Ubuntu20.04)+python3.8+pytorch1.14等等

2.1.SAM实现

2.1.0.环境设置

git clone https://github.com/facebookresearch/segment-anything
cd segment-anything; pip install -e .

pip install  pycocotools matplotlib onnxruntime onnx

# docker容器内的情况下务必安装  
pip install opencv-python-headless
pip install opencv-contrib-python 


pip install opencv-python==4.8.0.74


# b的大小:358M
wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth


# l的大小:1.2G
wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth

# 默认为h,h的大小:2.4G
wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth

2.1.1.分割方法实例

输入提示点分割

可以直接用官网给的,segment-anything/notebooks/predictor_example.ipynb文件。

可以用谷歌的,也也用自己的GPU。

import sys
sys.path.append("..")
from segment_anything import sam_model_registry, SamPredictor

# 这里指定sam预训练模型的权重地址及类型
sam_checkpoint = "sam_vit_h_4b8939.pth"
model_type = "vit_h"

# 指定cuda还是cpu
device = "cuda"

sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)

#用SamPredictor进行预测。该模型返回掩码、这些掩码的质量预测和低分辨率的掩码对数,可传递给下一次迭代预测。
predictor = SamPredictor(sam)

predictor.set_image(image)

输入点为提示,进行图像分割的情况

# 输入提示点的坐标[x,y]
input_point = np.array([[400,425]])
# 定义标签:[0]: 背景   [1]:前景
input_label = np.array([1])

masks, scores, logits = predictor.predict(
    point_coords=input_point,
    point_labels=input_label,
    multimask_output=True,
)

用`SamPredictor`进行预测。该模型返回掩码、这些掩码的质量预测和低分辨率的掩码对数,可传递给下一次迭代预测。

默认输出3个掩码可视化,使用multimask_output=True会生成3个mask,multimask_output=False的话只输出一个mask

全局分割

jupter lab的代码写到一起

# 导入相关库和优化可视化效果

import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2

def show_anns(anns):
    if len(anns) == 0:
        return
    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
    ax = plt.gca()
    ax.set_autoscale_on(False)
    polygons = []
    color = []
    for ann in sorted_anns:
        m = ann['segmentation']
        img = np.ones((m.shape[0], m.shape[1], 3))
        color_mask = np.random.random((1, 3)).tolist()[0]
        for i in range(3):
            img[:,:,i] = color_mask[i]
        ax.imshow(np.dstack((img, m*0.35)))

image = cv2.imread('test.jpg')
# 还原原图像色彩
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
plt.figure(figsize=(10,10))
plt.imshow(image)
plt.axis('on')
plt.show()


import sys
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor

# 加载模型
sam_checkpoint = '/lvdongrui/segment_anything/sam_vit_h_4b8939.pth'
model_type = "vit_h"

device = "cuda"

sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)

# 调用全局分割模型
mask_generator = SamAutomaticMaskGenerator(sam)

#预测并输出可视化结果
masks = mask_generator.generate(image)
print(len(masks)) # 产生的掩码数量
print(masks[0].keys()) # 第1个掩码内的相关属性

plt.figure(figsize=(10,10))
plt.imshow(image)
show_anns(masks)
plt.axis('off')
plt.show() 


# 如果想要调节相关参数的话
# 配置相关参数
mask_generator_2 = SamAutomaticMaskGenerator(
    model=sam,
    points_per_side=32,
    pred_iou_thresh=0.86, # IOU阈值
    stability_score_thresh=0.92, # 稳定性得分阈值
    crop_n_layers=1, 
    crop_n_points_downscale_factor=2,
    min_mask_region_area=100,  # Requires open-cv to run post-processing
)

masks2 = mask_generator_2.generate(image)
print(len(masks2)) # 产生的掩码数量
print(masks2[0].keys())

plt.figure(figsize=(10,10))
plt.imshow(image)
show_anns(masks2)
plt.axis('off')
plt.show() 

结果

2.1.2.基于SAM微调自己的数据集

数据集定义:dataloader.py



import numpy as np
import pandas as pd
import os
import cv2
import csv
import random

import torch
import torch.nn.functional as F
from torch.utils.data import Dataset
from torchvision import transforms
from sklearn.model_selection import train_test_split
from segment_anything import sam_model_registry
from segment_anything.utils.transforms import ResizeLongestSide


sam = sam_model_registry['vit_b'](checkpoint='./checkpoint/sam_vit_b_01ec64.pth')

class CustomDataset(Dataset):
    def __init__(self, data_dir, data_type="train"):
        self.data_dir = data_dir
        df = pd.read_csv(data_dir)

        train_names, test_names = train_test_split(df, test_size = 0.2, shuffle = True, random_state = 42) #  train, valid 
        val_names, test_names = train_test_split(test_names, test_size = 0.5, shuffle = True, random_state = 42) # valid, test 
        
        if data_type == "train":
            self.names = train_names
        if data_type == "valid":
            self.names = val_names
        if data_type == "test":
            self.names = test_names
            
        self.transform = ResizeLongestSide(1024)
        self.preprocess = sam.preprocess
        self.img_size = sam.image_encoder.img_size
        self.resize = transforms.Resize((256, 256))

    def __len__(self):
        
        return len(self.names)
    
    def __getitem__(self, idx):
        
        #####################################
        data = self.names

        image_path = data["image_path"].iloc[idx] # 读取image data路径 
        mask_path = data["mask_path"].iloc[idx] # 读取mask data 路径
        #####################################
        
        img = cv2.imread(image_path) # 读取原图数据
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        img = self.transform.apply_image(img) #
        img = torch.as_tensor(img) # torch tensor 变更
        img = img.permute(2, 0, 1).contiguous()[None, :, :, :].squeeze(0) # (高, 宽, 通道) -> (通道, 高, 宽) 变更后 设置添加None
        
        img = self.preprocess(img) # img nomalize or padding
        #####################################

        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) #读取掩码数据
        mask = self.transform.apply_image(mask) # 变换(1024)

        mask = torch.as_tensor(mask) # torch tensor 
        mask = mask.unsqueeze(0) 
        
        h, w = mask.shape[-2:] 
        
        padh = self.img_size - h 
        padw = self.img_size - w 
        
        mask = F.pad(mask, (0, padw, 0, padh)) 
        mask = self.resize(mask).squeeze(0) 
        mask = (mask != 0) * 1
        
        #####################################

        data = {
            'image': img, 
            'mask': mask, 
        }
        
        return data


 训练代码:train.py

# 导入了一些库
import warnings
warnings.filterwarnings(action='ignore')
import numpy as np
from tqdm import tqdm
import cv2
from datetime import datetime
import json
import os
import matplotlib.pyplot as plt
import argparse

import torch
from torch import nn, optim
import torch.nn.functional as F
from torch.nn.functional import threshold, normalize
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
from segment_anything.utils.transforms import ResizeLongestSide
from dataloader_sam import CustomDataset

# 设置了一些配置参数
beta = [0.9, 0.999]
milestone = [60000, 86666]
gamma = 0.1
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")

def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--batch_size', type=int, default=4, help = 'batch size')
    parser.add_argument('--warmup_steps', type = int, default = 250, help = ' ')
    parser.add_argument('--global_step', type = int, default = 0, help = ' ')
    parser.add_argument('--epochs', type = int, default = 10, help = 'train epcoh')
    parser.add_argument('--lr', type = float, default = 1e-5, help = 'learning_rate')
    parser.add_argument('--weight_decay', type = float, default = 0.1, help = 'weight_decay')
    parser.add_argument('--num_workers', type = int, default = 0, help = 'num_workers')
    parser.add_argument('--model_path', type = str, default = './models/', help = 'model path directory')
    parser.add_argument('--data_dir', type = str, default = './data_list_sam/pr_data_list.csv', help = 'data directory')
    parser.add_argument('--pretrained', type = str, default = False, help = 'pre trained model select')
    parser.add_argument('--device_id', type = int, default = 0, help = 'Cuda device Id')
    

    return parser.parse_known_args()[0]

# 数据加载
def build_dataloader(data_dir, batch_size, num_workers):

    dataloaders = {
        key: DataLoader(
            CustomDataset(data_dir, key),
            batch_size = batch_size,
            shuffle = True if key != 'test' else False,
            num_workers = num_workers,
            pin_memory = True
        ) for key in ['train', 'valid', 'test']
    }

    return dataloaders

# 损失函数
def focal_loss(pred, target, gamma=2.0, alpha=0.25, reduction='mean'):
    #pred = F.sigmoid(pred)
    pt = torch.where(target == 1, pred, 1-pred)
    ce_loss = F.binary_cross_entropy(pred, target, reduction="none")
    focal_term = (1 - pt).pow(gamma)
    loss = alpha * focal_term * ce_loss

    return loss.mean()

def dice_loss(pred, target, smooth=1.0):
    pred_flat = pred.reshape(-1)
    target_flat = target.reshape(-1)
    intersection = (pred_flat * target_flat).sum()

    return 1 - ((2. * intersection + smooth) /
              (pred_flat.sum() + target_flat.sum() + smooth))

def compute_loss(pred_mask, true_mask, pred_iou, true_iou):
    pred_mask = F.sigmoid(pred_mask).squeeze(1).to(dtype = torch.float64)
    fl = focal_loss(pred_mask, true_mask)
    dl = dice_loss(pred_mask, true_mask)
    mask_loss = 20 * fl + dl
    iou_loss = F.mse_loss(pred_iou, true_iou)
    total_loss = mask_loss + iou_loss

    return total_loss

def mean_iou(preds, labels, eps=1e-6):
    preds = normalize(threshold(preds, 0.0, 0)).squeeze(1)
    pred_cls = (preds == 1).float()
    label_cls = (labels == 1).float()
    intersection = (pred_cls * label_cls).sum(1).sum(1)
    union = (1 - (1 - pred_cls) * (1 - label_cls)).sum(1).sum(1)
    intersection = intersection + (union == 0)
    union = union + (union == 0)
    ious = intersection / union

    return ious


def main(opt):
    use_cuda = torch.cuda.is_available()
    device = torch.device(f"cuda:{opt.device_id}" if use_cuda else "cpu")
    torch.manual_seed(234)
    if device == f'cuda:{opt.device_id}':
        torch.cuda.manual_seed_all(234)
    #  脚本使用预先构建的架构(sam_model_registry['vit_b'])定义了一个神经网络模型,并设置了优化器(AdamW)和学习率调度。
    print(device,'is available')
    print("Loading model...")
    epoch_add = 0
    lr = opt.lr
    sam = sam_model_registry['vit_b'](checkpoint='./checkpoint/sam_vit_b_01ec64.pth')


    if opt.pretrained:
        sam.load_state_dict(torch.load('./models/' + opt.pretrained))
        sam = sam.to(device=device)
    else: 
        sam = sam.to(device=device)
        sam_decoder = sam.mask_decoder.to(device=device)
        sam_decoder = sam.mask_decoder

    optimizer = optim.AdamW(sam.mask_decoder.parameters(),
                           lr = lr, betas = beta, weight_decay = opt.weight_decay)
    criterion = nn.CrossEntropyLoss()
    scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=milestone, gamma=gamma)


    # 脚本在各个检查点保存训练模型的状态字典,如果模型在验证集上取得最佳平均IOU,则单独保存最佳模型。
    if len(os.listdir(opt.model_path)) == 0:
        save_path = os.path.join(opt.model_path, f"model_{opt.epochs}_{opt.batch_size}_0")
        os.makedirs(save_path)
    else:
        save_path = os.path.join(opt.model_path, f"model_{opt.epochs}_{opt.batch_size}_" + str(len(os.listdir(opt.model_path))))
        os.makedirs(save_path)

    print('Training Start')
    best_loss = 999999999

    best_mIOU = 0

    tr_pl_loss_list = []
    tr_pl_mi_list = []
    val_pl_loss_list = []
    val_pl_mi_list = []

    dataloaders = build_dataloader(opt.data_dir, opt.batch_size, opt.num_workers)
    for epoch in range(opt.epochs):
        train_loss_list = []
        train_miou_list = []

        sam.train()
        iterations = tqdm(dataloaders['train'])

        # 循环进行模型的多轮训练
        for train_data in iterations:
            # 将训练数据移到指定设备,这里是GPU
            train_input = train_data['image'].to(device)
            train_target_mask = train_data['mask'].to(device, dtype = torch.float64)
            # 对优化器的梯度进行归零
            optimizer.zero_grad()

            with torch.no_grad():
                # 使用 sam 模型的 image_encoder 提取图像特征,并使用 prompt_encoder 提取稀疏和密集的嵌入。在本代码中进行提示输入,所以都是None.
                train_encode_feature = sam.image_encoder(train_input)
                train_sparse_embeddings, train_dense_embeddings = sam.prompt_encoder(points = None, boxes = None, masks = None)
            
              #  通过 mask_decoder 解码器生成训练集的预测掩码和IOU
            train_mask, train_IOU = sam.mask_decoder(

image_embeddings = train_encode_feature,
                                                     
image_pe = sam.prompt_encoder.get_dense_pe(),
                                                     
sparse_prompt_embeddings = train_sparse_embeddings,
                                                     
dense_prompt_embeddings = train_dense_embeddings,
                                                     
multimask_output = False)
            # 计算预测IOU和真实IOU之间的差异,并将其添加到列表中。然后计算训练损失(总损失包括mask损失和IOU损失),进行反向传播和优化器更新。
            train_true_iou = mean_iou(train_mask, train_target_mask, eps=1e-6)
            train_miou_list = train_miou_list + train_true_iou.tolist()

            train_loss_one = compute_loss(train_mask, train_target_mask, train_IOU, train_true_iou)
            train_loss_one.backward()

            optimizer.step()
 
            train_loss_list.append(train_loss_one.item())     
            # 学习率调整
            if epoch_add == 0:

                if opt.global_step < opt.warmup_steps:    
                    lr_scale = opt.global_step / opt.warmup_steps
                    for param_group in optimizer.param_groups:
                        param_group['lr'] = 8e-4 * lr_scale

                opt.global_step += 1


            pbar_desc = "Model train loss --- "
            pbar_desc += f"Total loss: {np.mean(train_loss_list):.5f}"
            pbar_desc += f", total mIOU: {np.mean(train_miou_list):.5f}"
            iterations.set_description(pbar_desc)


        train_loss = np.mean(train_loss_list)
        train_miou = np.mean(train_miou_list)

        torch.cuda.empty_cache()
        tr_pl_loss_list.append(train_loss)
        tr_pl_mi_list.append(train_miou)

        sam.eval()
        scheduler.step()

        with torch.no_grad():
            valid_loss_list = []
            valid_miou_list = []
            valid_true_iou = 0
            valid_loss = 0
            valid_miou = 0

            iterations = tqdm(dataloaders['valid'])

            for valid_data in iterations:   
                valid_input = valid_data['image'].to(device)
                valid_target_mask = valid_data['mask'].to(device, dtype = torch.float64)
            

                valid_encode_feature = sam.image_encoder(valid_input)
                valid_sparse_embeddings, valid_dense_embeddings = sam.prompt_encoder(points = None, boxes = None, masks = None)


                valid_mask, valid_IOU = sam.mask_decoder(image_embeddings = valid_encode_feature,

                                                        image_pe = sam.prompt_encoder.get_dense_pe(),

                                                        sparse_prompt_embeddings = valid_sparse_embeddings,

                                                        dense_prompt_embeddings = valid_dense_embeddings,

                                                        multimask_output = False)

                valid_true_iou = mean_iou(valid_mask, valid_target_mask, eps=1e-6)
                valid_miou_list = valid_miou_list + valid_true_iou.tolist()

                valid_loss_one = compute_loss(valid_mask, valid_target_mask, valid_IOU, valid_true_iou)
                valid_loss_list.append(valid_loss_one.item())

                pbar_desc = "Model valid loss --- "
                pbar_desc += f"Total loss: {np.mean(valid_loss_list):.5f}"
                pbar_desc += f", total mIOU: {np.mean(valid_miou_list):.5f}"
                iterations.set_description(pbar_desc)


            valid_loss = np.mean(valid_loss_list)
            valid_miou = np.mean(valid_miou_list)
            val_pl_loss_list.append(valid_loss)
            val_pl_mi_list.append(valid_miou)       


        model_path = opt.model_path + 'sam.pth'
        sam = sam.cpu()
        torch.save(sam.state_dict(), model_path)
        sam = sam.to(device)


        if best_mIOU < valid_miou:
            best_loss = valid_loss
            best_mIOU = valid_miou
            model_path = save_path + '/sam_best.pth'
            sam = sam.cpu()
            torch.save(sam.state_dict(), model_path)
            sam = sam.to(device)
            f = open(os.path.join(save_path, 'best.txt'), 'w')
            f.write(f"Experimental Day: {datetime.now()}")
            f.write("\n")
            f.write(f"mIoU: {str(best_mIOU)}")
            f.write("\n")
            f.write(f"epochs:{opt.epochs}")
            f.write("\n")
            f.write(f"batch_size:{opt.batch_size}")
            f.write("\n")
            f.write(f"learning_rate:{opt.lr}")
            f.write("\n")
            f.write(f"Data_set : {opt.data_dir}")
            f.close()

        print("epoch : {:3d}, train loss : {:3.4f}, valid loss : {:3.4f}, valid mIOU : {:3.4f}\
            ( best vaild loss : {:3.4f}, best valid mIOU : {:3.4f} )".format(epoch + 1 + epoch_add,
                                                                             train_loss,
                                                                             valid_loss,
                                                                             valid_miou,
                                                                             best_loss,
                                                                             best_mIOU
                                                                             ))

        lr = optimizer.param_groups[0]["lr"]

        if (epoch + 1) % 5 == 0 or (epoch + 1) in [1, 2, 3, 4, 5]:
            model_path = save_path + "/"+ "sam_" + str(epoch + 1 + epoch_add) + '_' + str(round(lr, 10)) + '.pth'
            sam = sam.cpu()
            torch.save(sam.state_dict(), model_path)
            sam = sam.to(device)

    # (2, 2) 形式的图使用matplotlib可视化训练进展,生成用于训练和验证平均IOU、训练和验证损失的图表。
    plt_dict = {
        "Train_mIoU" : tr_pl_mi_list,
        "Val_mIoU" : val_pl_mi_list,
        "Train_Loss" : tr_pl_loss_list,
        "Val_Loss" : val_pl_loss_list
    }

    plt.figure(figsize=(15,15))
    for i, (key, item) in enumerate(plt_dict.items()):
        plt.subplot(2, 2, i+1)
        plt.plot(range(opt.epochs), item, label=f"{key}")
        plt.title(f"{key}", fontsize = 20)
        plt.xlabel('Epochs', fontsize = 15)
        plt.ylabel(f'{key.split("_")[-1]}', fontsize = 15)
        plt.grid(True)

    plt.savefig(save_path + f'/sam_{opt.epochs}_{opt.batch_size}_{opt.lr}_result.png')

if __name__ == '__main__':
    opt = parse_opt()
    main(opt)

2.2.YOLO+SAM

2.2.0.环境设置及Demo

pip install ultralytics
#pip install git+https://github.com/facebookresearch/segment-anything.git
#wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth

git clone https://github.com/poojatambe/Segmentation_models_on_custom_data
cd Segmentation_models_on_custom_data

# 下载yolov8的检测权重
wget https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8x-oiv7.pt

# 根据自己路径修改
python SAM_yolov8.py --img './test.png' --weights './yolov8-oiv7.pt' --model_type 'vit_l' --checkpoints './sam_vit_l_0b3195.pth'

如果出现错误,请参考【PS3】

3.2.1.微调自己的数据集

敬请期待。。。

3.相关知识点进阶

输入提示点实例

#输入1个点

input_point = np.array([[495,518]])
input_label = np.ones(input_point.shape[0])
input_box = None

#输入三个提示点

input_point = np.array([[221,482],[498,633],[750,379]])
input_label = np.ones(input_point.shape[0])
input_box = None

输入提示框box实例

#输入一个box

input_box = np.array([[4,13,1007,1023]])
input_point, input_label = None, None

#输入多个box

input_box = torch.tensor([[45,260,515,470], [310,228,424,296]],device=predictor.device)
transformed_box = predictor.transform.apply_boxes_torch(input_box, image.shape[:2])
input_point, input_label = None, None

 使用box的计算

IoU(Intersection over Union): 交并比,衡量bounding box与ground truth的重叠程度。比值越大重叠度越高,完全重合为1。如果还没有bbox回归,则bbox就是候选框/锚框。

IoU= \frac{boundingbox\cap groundtruth}{boundingbox\cup groundtruth}

通常表示边界框的位置有两种方式:

  • (x1​,y1​,x2​,y2​):其中(x1​,y1​)是矩形框左上角的坐标,(x2​,y2​)是矩形框右下角的坐标。
  • (x,y,w,h):其中(x,y)是矩形框中心点的坐标,w是矩形框的宽度,h是矩形框的高度。

在SAM中使用xywh 格式的掩模边界框

 【PS1】IndexError: boolean index did not match indexed array along dimension 0; dimension is 2 but corresponding boolean dimension is 1

错误原因:数组输入错误,少了一对中括号

将自定义的input_point = np.array([200,400])
改为input_point = np.array([[200,400]])

【PS2】RuntimeError:CUDNN_BACKEND_TENSOR_DESCRIPTOR: Check and Set the CUDNN_ATTR_TENSOR_STRIDES Correctly cudnn_status: CUDNN_STATUS_BAD_PARAM

【PS3】RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument mat2 in method wrapper_mm)

这是由于代码所写的某一位置的Tensor设备指代不清楚,tensors一会在cpu计算,一会在GPU计算。容易导致该错误产生的原因:

当多GPU并行训练时(我的是俩块):

这时候你的model里面传入的不只是Cuda:0,还有Cuda:1, Cuda:2等等,这个时候,你的网络模型model里面的forward函数里面就及其容易报错,因为里面存在的一些定维度的Tensor,比如权重、偏差的device容易跑偏。

当单GPU正常使用Cuda:0:

这时候按理来说,只要device代码指代正常,按理来说不会出现设备问题,如果出现了,就是tensor的格式没有传入cuda而是保留在cpu。碰巧我某天训练时候就出现了,思路跟上面一样,最容易出现错误的地方是网络模型中固定维度的tensor如权重和偏差等tensor,或者是forward函数里面的输出.

解决

指定GPU

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

 更多可参考RuntimeError: Expected all tensors to be on the same device, but found at least two devices-CSDN博客

Q&A

【Q&A1】怎么计算mask的shape?

在图像分割中,模型在训练过程中计算mask的形状(shape)取决于任务的要求。分割模型通常输出一个与输入图像相同空间分辨率的mask,该mask用于标识图像中的不同区域或对象。这个mask的形状通常与输入图像的宽度和高度相匹配。

参考文献

【1】SAM : Segment Anything - Best Simple Tutorial to master it (inside-machinelearning.com)

【2】SAM + YOLO8图像检测与分割简明教程 - 知乎 (zhihu.com)

【3】Segment Anything Model (SAM): Explained | by Utkarsh Doshi | Dec, 2023 | Medium 

【4】SAM: A Image Segmentation Foundation Model – Towards AI 

【5】SAM(segment anything model)分割一切 Demo测试及API调用_sam模型运行实例-CSDN博客 

【6】Segment Anything(SAM)的demo的简单使用_segment anything demo-CSDN博客 

  • 28
    点赞
  • 35
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 2
    评论
Segment Anything Model可以指很多不同的模型,因此调用方式也不同。以下是使用Mask R-CNN进行实例分割的示例代码: ``` import os import sys import random import math import numpy as np import skimage.io import matplotlib import matplotlib.pyplot as plt # Root directory of the project ROOT_DIR = os.path.abspath("../") # Import Mask RCNN sys.path.append(ROOT_DIR) # To find local version of the library from mrcnn.config import Config from mrcnn import model as modellib, utils # Directory to save logs and trained model MODEL_DIR = os.path.join(ROOT_DIR, "logs") # Path to trained weights file COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5") # Download COCO trained weights from Releases if needed if not os.path.exists(COCO_MODEL_PATH): utils.download_trained_weights(COCO_MODEL_PATH) # Directory of images to run detection on IMAGE_DIR = os.path.join(ROOT_DIR, "images") class InferenceConfig(Config): # Set batch size to 1 since we'll be running inference on # one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU GPU_COUNT = 1 IMAGES_PER_GPU = 1 NAME = "segment_anything_model" # Number of classes (including background) NUM_CLASSES = 1 + 80 # COCO has 80 classes config = InferenceConfig() config.display() # Create model object in inference mode. model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config) # Load weights trained on MS-COCO model.load_weights(COCO_MODEL_PATH, by_name=True) # COCO Class names # Index of the class in the list is its ID. For example, to get ID of # the teddy bear class, use: class_names.index('teddy bear') class_names = ['BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'] # Load a random image from the images folder file_names = next(os.walk(IMAGE_DIR))[2] image = skimage.io.imread(os.path.join(IMAGE_DIR, random.choice(file_names))) # Run detection results = model.detect([image], verbose=1) # Visualize results r = results[0] visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'], class_names, r['scores']) ``` 这段代码会在指定的文件夹中随机选取一张图片进行实例分割,并将结果可视化显示。你需要将代码中的`ROOT_DIR`、`MODEL_DIR`、`COCO_MODEL_PATH`、`IMAGE_DIR`等路径修改为你自己的路径。同时,如果你使用的是其他的Segment Anything Model,那么需要根据具体的模型进行修改。

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 2
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

夏天|여름이다

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值