取得图像后，直接使用SamPredictor中的预处理方式，会将图片按最长边resized到1024x1024，然后计算image_embedding，这部分很耗时，所以每张图只计算一次，会将结果缓存起来需要的时候直接调用。使用"with torch.no_grad()"保证image encoder部分不需要梯度更新，冻结对应权重

model_transform = ResizeLongestSide(sam.image_encoder.img_size)

for epoch in range(num_epochs):

epoch_loss = 0

for idx, (images, gts, image_names) in enumerate(tqdm(dataloader)):

valid_classes = [] ## voc 0,255 are ignored

for i in range(images.shape[0]):

image = images[i] # h,w,c np.uint8 rgb

original_size = image.shape[:2] ## h,w

input_size = model_transform.get_preprocess_shape(image.shape[0], image.shape[1],

sam.image_encoder.img_size) ##h,w

gt = gts[i].copy() #h,w labels [0,1,2,..., classes-1]

gt_classes = np.unique(gt) ##masks classes: [0, 1, 2, 3, 4, 7]

image_name = image_names[i]

predictions = []

## freeze image encoder

with torch.no_grad():

# gt_channel = gt[:, :, cls]

predictor.set_image(image, "RGB")

image_embedding = predictor.get_image_embedding()

Prompt生成

从mask中随机选取一定数量的前景点和背景点，此处默认1个前景点和1个背景点，数量多的话一般保持2:1的比例较好。

mask_value就是对应的类别id，去mask中找出像素值等于类别id的点坐标，然后随机选取点就行。此处还会根据mask算外接矩形（实际上直接读取图片对应的xml标签文件也行），用于后续基于box prompt的finetune。

def get_random_prompts(mask, mask_value, foreground_nums=1, background_nums=1):

# Find the indices (coordinates) of the foreground pixels

foreground_indices = np.argwhere(mask == mask_value)

ymin, xmin= foreground_indices.min(axis=0)

ymax, xmax = foreground_indices.max(axis=0)

bbox = np.array([xmin, ymin, xmax, ymax])

if foreground_indices.shape[0] < foreground_nums:

foreground_nums = foreground_indices.shape[0]

background_nums = int(0.5 * foreground_indices.shape[0])

background_indices = np.argwhere(mask != mask_value)

## random select

foreground_points = foreground_indices[

np.random.choice(foreground_indices.shape[0], foreground_nums, replace=False)]

background_points = background_indices[

np.random.choice(background_indices.shape[0], background_nums, replace=False)]

## 坐标点是(y,x)，输入给网络应该是(x,y),需要翻一下顺序

foreground_points = foreground_points[:, ::-1]

background_points = background_points[:, ::-1]

return (foreground_points, background_points), bbox

得到的prompt是一些点的坐标，坐标的x,y是基于原图的，但进入SAM的图片会resized到1024x1024，所以点坐标也需要resize，对应如下代码

all_points = np.concatenate((foreground_points, background_points), axis=0)

all_points = np.array(all_points)

point_labels = np.array([1] * foreground_points.shape[0] + [0] * background_points.shape[0], dtype=int)

## image resized to 1024, points also

all_points = model_transform.apply_coords(all_points, original_size)

all_points = torch.as_tensor(all_points, dtype=torch.float, device=device)

point_labels = torch.as_tensor(point_labels, dtype=torch.float, device=device)

all_points, point_labels = all_points[None, :, :], point_labels[None, :]

points = (all_points, point_labels)

if not box_prompt:

box_torch=None

else:

## preprocess bbox

box = model_transform.apply_boxes(bbox, original_size)

box_torch = torch.as_tensor(box, dtype=torch.float, device=device)

box_torch = box_torch[None, :]

微调代码中可以指定基于哪种prompt进行微调，如果是point和box同时都开，会按一定概率舍弃point或box以取得更好的泛化性（不然推理时只有point或只有box作为prompt效果可能不太好）。最后经过prompt_encoder得到sparse_embeddings, dense_embeddings。

## if both, random drop one for better generalization ability

if point_box and np.random.random()<0.5:

if np.random.random()<0.25:

points = None

elif np.random.random()>0.75:

box_torch = None

## freeze prompt encoder

with torch.no_grad():

sparse_embeddings, dense_embeddings = sam.prompt_encoder(

points = points,

boxes = box_torch,

# masks=mask_predictions,

masks=None,

)

Mask预测

mask decoder这部分不需要冻结，直接调用mask_decoder推理就行，这里进行了两次mask预测，第一次先预测3个层级的mask然后选出得分最高的一个，将这个mask作为一个mask prompt，并与point prompt、box_prompt一起丢进prompt_encoder得到新的sparse_embeddings, dense_embeddings，再进行第二次mask预测，这次只预测一个mask。相当于先得到粗糙的mask，然后再精修。最后经过后处理nms等得到和原图大小一样的预测mask，一个物体对应一张mask，将多个mask叠起来就得到这张图所有的预测结果predictions。

## predicted masks, three level

mask_predictions, scores = sam.mask_decoder(

image_embeddings=image_embedding.to(device),

image_pe=sam.prompt_encoder.get_dense_pe(),

sparse_prompt_embeddings=sparse_embeddings,

dense_prompt_embeddings=dense_embeddings,

multimask_output=True,

)

# Choose the model's best mask

mask_input = mask_predictions[:, torch.argmax(scores),...].unsqueeze(1)

with torch.no_grad():

sparse_embeddings, dense_embeddings = sam.prompt_encoder(

points=points,

boxes=box_torch,

masks=mask_input,

)

## predict a better mask, only one mask

mask_predictions, scores = sam.mask_decoder(

image_embeddings=image_embedding.to(device),

image_pe=sam.prompt_encoder.get_dense_pe(),

sparse_prompt_embeddings=sparse_embeddings,

dense_prompt_embeddings=dense_embeddings,

multimask_output=False,

)

best_mask = sam.postprocess_masks(mask_predictions, input_size, original_size)

predictions.append(best_mask)

Loss计算

代码中loss用的是BCELoss加DiceLoss，需要gt和pred的shape一致，都为BxCxHxW的形式，pred是经过sigmoid后的值。

因此需要将gt转为one-hot的形式，即将(batch_size, 1, h, w)转为(batch_size, c, h, w)，c是gt_classes中有的类别个数，即图片中有多少个实例类别。

def mask2one_hot(label, gt_classes):

"""

label: 标签图像 # (batch_size, 1, h, w)

num_classes: 分类类别数

"""

current_label = label.squeeze(1) # （batch_size, 1, h, w) ---> （batch_size, h, w)

batch_size, h, w = current_label.shape[0], current_label.shape[1], current_label.shape[2]

one_hots = []

for cls in gt_classes:

if isinstance(cls, torch.Tensor):

cls = cls.item()

tmplate = torch.zeros(batch_size, h, w) # （batch_size, h, w)

tmplate[current_label == cls] = 1

tmplate = tmplate.view(batch_size, 1, h, w) # （batch_size, h, w) --> （batch_size, 1, h, w)

one_hots.append(tmplate)

onehot = torch.cat(one_hots, dim=1)

return onehot

另外BCE接受的pred值是logit形式，所以需要将predictions用sigmoid处理，后续loss计算对应如下代码

gts = torch.from_numpy(gts).unsqueeze(1) ## BxHxW ---> Bx1xHxW

gts_onehot = mask2one_hot(gts, valid_classes)

gts_onehot = gts_onehot.to(device)

predictions = torch.sigmoid(predictions)

# #loss = seg_loss(predictions, gts_onehot)

loss = BCEseg(predictions, gts_onehot)

loss_dice = soft_dice_loss(predictions, gts_onehot, smooth = 1e-5, activation='none')code>

loss = loss + loss_dice

权重保存