341 lines
13 KiB
Python
341 lines
13 KiB
Python
from typing import List
|
|
import os
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from transformers import PreTrainedModel, AutoConfig, AutoModelForCausalLM
|
|
from .segment_anything_2.sam2.build_sam import build_sam2
|
|
from .unilm.beit3.modeling_utils import BEiT3Wrapper, _get_base_config, _get_large_config
|
|
from .configuration_evf import EvfConfig
|
|
from .segment_anything_2.sam2.utils.misc import load_video_frames
|
|
from collections import OrderedDict
|
|
|
|
|
|
def dice_loss(
|
|
inputs: torch.Tensor,
|
|
targets: torch.Tensor,
|
|
num_masks: float,
|
|
scale=1000, # 100000.0,
|
|
eps=1e-6,
|
|
):
|
|
"""
|
|
Compute the DICE loss, similar to generalized IOU for masks
|
|
Args:
|
|
inputs: A float tensor of arbitrary shape.
|
|
The predictions for each example.
|
|
targets: A float tensor with the same shape as inputs. Stores the binary
|
|
classification label for each element in inputs
|
|
(0 for the negative class and 1 for the positive class).
|
|
"""
|
|
inputs = inputs.sigmoid()
|
|
inputs = inputs.flatten(1, 2)
|
|
targets = targets.flatten(1, 2)
|
|
numerator = 2 * (inputs / scale * targets).sum(-1)
|
|
denominator = (inputs / scale).sum(-1) + (targets / scale).sum(-1)
|
|
loss = 1 - (numerator + eps) / (denominator + eps)
|
|
loss = loss.sum() / (num_masks + 1e-8)
|
|
return loss
|
|
|
|
|
|
def sigmoid_ce_loss(
|
|
inputs: torch.Tensor,
|
|
targets: torch.Tensor,
|
|
num_masks: float,
|
|
):
|
|
"""
|
|
Args:
|
|
inputs: A float tensor of arbitrary shape.
|
|
The predictions for each example.
|
|
targets: A float tensor with the same shape as inputs. Stores the binary
|
|
classification label for each element in inputs
|
|
(0 for the negative class and 1 for the positive class).
|
|
Returns:
|
|
Loss tensor
|
|
"""
|
|
loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none")
|
|
loss = loss.flatten(1, 2).mean(1).sum() / (num_masks + 1e-8)
|
|
return loss
|
|
|
|
class EvfSam2Model(PreTrainedModel):
|
|
config_class = EvfConfig
|
|
def __init__(
|
|
self,
|
|
config,
|
|
**kwargs
|
|
):
|
|
super(EvfSam2Model, self).__init__(config)
|
|
|
|
self.config = config
|
|
self.vision_pretrained = kwargs.get("vision_pretrained", None)
|
|
self.encoder_pretrained = kwargs.get("encoder_pretrained", None)
|
|
self.dice_loss_weight = kwargs.get("dice_loss_weight", None)
|
|
self.bce_loss_weight = kwargs.get("bce_loss_weight", None)
|
|
self.train_mask_decoder = kwargs.get("train_mask_decoder", False)
|
|
self.train_prompt_encoder = kwargs.get("train_prompt_encoder", False)
|
|
self.initialize_evf_modules(config)
|
|
self._bb_feat_sizes = [
|
|
(256, 256),
|
|
(128, 128),
|
|
(64, 64),
|
|
]
|
|
|
|
def initialize_evf_modules(self, config):
|
|
# SAM
|
|
if config.sam_scale=="large":
|
|
self.visual_model = build_sam2("sam2_hiera_l.yaml", self.vision_pretrained, device=None)
|
|
elif config.sam_scale=="tiny":
|
|
self.visual_model = build_sam2("sam2_hiera_t.yaml", self.vision_pretrained, device=None)
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
for param in self.visual_model.parameters():
|
|
param.requires_grad = False
|
|
if self.train_mask_decoder:
|
|
self.visual_model.sam_mask_decoder.train()
|
|
for param in self.visual_model.sam_mask_decoder.parameters():
|
|
param.requires_grad = True
|
|
if self.train_prompt_encoder:
|
|
self.visual_model.sam_prompt_encoder.no_mask_embed.requires_grad_(True)
|
|
|
|
# beit-3
|
|
if self.config.mm_extractor_scale == "base":
|
|
beit_config = _get_base_config()
|
|
elif self.config.mm_extractor_scale == "large":
|
|
beit_config = _get_large_config()
|
|
else:
|
|
raise AttributeError(f"model config should contain key 'mm_extractor_scale', with value 'base' or 'large'.")
|
|
|
|
self.mm_extractor = BEiT3Wrapper(beit_config)
|
|
if self.encoder_pretrained is not None:
|
|
beit_state_dict = torch.load(self.encoder_pretrained)["model"]
|
|
self.mm_extractor.load_state_dict(
|
|
beit_state_dict,
|
|
strict=False
|
|
)
|
|
|
|
for param in self.mm_extractor.parameters():
|
|
param.requires_grad = True
|
|
|
|
# Projection layer
|
|
in_dim = config.hidden_size
|
|
assert in_dim==beit_config.encoder_embed_dim, \
|
|
f"projection layer dim {in_dim} mismatch with mm_extractor dim {beit_config.encoder_embed_dim}"
|
|
out_dim = config.out_dim
|
|
text_fc = [
|
|
nn.Linear(in_dim, in_dim),
|
|
nn.ReLU(),
|
|
nn.Linear(in_dim, out_dim)
|
|
]
|
|
self.text_hidden_fcs = nn.ModuleList([nn.Sequential(*text_fc)])
|
|
self.text_hidden_fcs.train()
|
|
for param in self.text_hidden_fcs.parameters():
|
|
param.requires_grad = True
|
|
|
|
def postprocess_masks(self, masks: torch.Tensor, orig_hw) -> torch.Tensor:
|
|
"""
|
|
Perform PostProcessing on output masks.
|
|
"""
|
|
masks = masks.float()
|
|
masks = F.interpolate(masks, orig_hw, mode="bilinear", align_corners=False)
|
|
return masks
|
|
|
|
def forward(
|
|
self,
|
|
images: torch.FloatTensor,
|
|
images_evf: torch.FloatTensor,
|
|
input_ids: torch.LongTensor,
|
|
attention_masks: torch.LongTensor,
|
|
offset: torch.LongTensor,
|
|
masks_list: List[torch.FloatTensor],
|
|
label_list: List[torch.Tensor],
|
|
resize_list: List[tuple],
|
|
inference: bool = False,
|
|
**kwargs,
|
|
):
|
|
# image_embeddings = self.get_visual_embs(images)
|
|
backbone_out = self.visual_model.forward_image(images)
|
|
# dict_keys(['vision_features', 'vision_pos_enc', 'backbone_fpn'])
|
|
_, image_embeddings, _, _ = self.visual_model._prepare_backbone_features(backbone_out)
|
|
image_embeddings = [_.to(images.dtype) for _ in image_embeddings]
|
|
batch_size = images.shape[0]
|
|
if self.visual_model.directly_add_no_mem_embed:
|
|
image_embeddings[-1] = image_embeddings[-1] + self.visual_model.no_mem_embed
|
|
|
|
feats = [
|
|
feat.permute(1, 2, 0).view(batch_size, -1, *feat_size)
|
|
for feat, feat_size in zip(image_embeddings[::-1], self._bb_feat_sizes[::-1])
|
|
][::-1]
|
|
_features = {"image_embed": feats[-1], "high_res_feats": feats[:-1]}
|
|
|
|
|
|
assert batch_size == len(offset) - 1
|
|
|
|
images_evf_list = []
|
|
for i in range(len(offset) - 1):
|
|
start_i, end_i = offset[i], offset[i + 1]
|
|
images_evf_i = (
|
|
images_evf[i]
|
|
.unsqueeze(0)
|
|
.expand(end_i - start_i, -1, -1, -1)
|
|
.contiguous()
|
|
)
|
|
images_evf_list.append(images_evf_i)
|
|
images_evf = torch.cat(images_evf_list, dim=0)
|
|
|
|
multimask_output = False
|
|
output = self.mm_extractor.beit3(
|
|
visual_tokens=images_evf,
|
|
textual_tokens=input_ids,
|
|
text_padding_position=~attention_masks
|
|
)
|
|
|
|
feat = output["encoder_out"][:, :1, ...]
|
|
|
|
feat = self.text_hidden_fcs[0](feat)
|
|
feat = torch.split(feat, [offset[i+1] - offset[i] for i in range(len(offset)-1)])
|
|
|
|
pred_masks = []
|
|
|
|
for i in range(len(feat)):
|
|
(
|
|
sparse_embeddings,
|
|
dense_embeddings,
|
|
) = self.visual_model.sam_prompt_encoder(
|
|
points=None,
|
|
boxes=None,
|
|
masks=None,
|
|
text_embeds=feat[i],
|
|
)
|
|
sparse_embeddings = sparse_embeddings.to(feat[i].dtype)
|
|
high_res_features = [
|
|
feat_level[i].unsqueeze(0)
|
|
for feat_level in _features["high_res_feats"]
|
|
]
|
|
low_res_masks, iou_predictions, _, _ = self.visual_model.sam_mask_decoder(
|
|
image_embeddings=_features["image_embed"][i].unsqueeze(0),
|
|
image_pe=self.visual_model.sam_prompt_encoder.get_dense_pe(),
|
|
sparse_prompt_embeddings=sparse_embeddings,
|
|
dense_prompt_embeddings=dense_embeddings,
|
|
multimask_output=multimask_output,
|
|
repeat_image = True,
|
|
high_res_features=high_res_features,
|
|
)
|
|
|
|
if multimask_output:
|
|
sorted_ids = torch.argsort(iou_predictions, dim=-1, descending=True)
|
|
low_res_masks = torch.take_along_dim(low_res_masks, sorted_ids[..., None, None], dim=1)[:, :1]
|
|
|
|
pred_mask = self.postprocess_masks(
|
|
low_res_masks,
|
|
orig_hw=label_list[i].shape,
|
|
)
|
|
pred_masks.append(pred_mask[:, 0])
|
|
|
|
gt_masks = masks_list
|
|
|
|
if inference:
|
|
return {
|
|
"pred_masks": pred_masks,
|
|
"gt_masks": gt_masks,
|
|
}
|
|
|
|
mask_bce_loss = 0
|
|
mask_dice_loss = 0
|
|
num_masks = 0
|
|
for batch_idx in range(len(pred_masks)):
|
|
gt_mask = gt_masks[batch_idx]
|
|
pred_mask = pred_masks[batch_idx]
|
|
|
|
assert (
|
|
gt_mask.shape[0] == pred_mask.shape[0]
|
|
), "gt_mask.shape: {}, pred_mask.shape: {}".format(
|
|
gt_mask.shape, pred_mask.shape
|
|
)
|
|
mask_bce_loss += (
|
|
sigmoid_ce_loss(pred_mask, gt_mask, num_masks=gt_mask.shape[0])
|
|
* gt_mask.shape[0]
|
|
)
|
|
mask_dice_loss += (
|
|
dice_loss(pred_mask, gt_mask, num_masks=gt_mask.shape[0])
|
|
* gt_mask.shape[0]
|
|
)
|
|
num_masks += gt_mask.shape[0]
|
|
|
|
mask_bce_loss = self.bce_loss_weight * mask_bce_loss / (num_masks + 1e-8)
|
|
mask_dice_loss = self.dice_loss_weight * mask_dice_loss / (num_masks + 1e-8)
|
|
mask_loss = mask_bce_loss + mask_dice_loss
|
|
|
|
loss = mask_loss
|
|
|
|
return {
|
|
"loss": loss,
|
|
"mask_bce_loss": mask_bce_loss,
|
|
"mask_dice_loss": mask_dice_loss,
|
|
"mask_loss": mask_loss,
|
|
}
|
|
|
|
def inference(
|
|
self,
|
|
images,
|
|
images_evf,
|
|
input_ids,
|
|
resize_list,
|
|
original_size_list,
|
|
multimask_output=False,
|
|
):
|
|
with torch.no_grad():
|
|
backbone_out = self.visual_model.forward_image(images)
|
|
# dict_keys(['vision_features', 'vision_pos_enc', 'backbone_fpn'])
|
|
_, image_embeddings, _, _ = self.visual_model._prepare_backbone_features(backbone_out)
|
|
image_embeddings = [_.to(images.dtype) for _ in image_embeddings]
|
|
batch_size = images.shape[0]
|
|
if self.visual_model.directly_add_no_mem_embed:
|
|
image_embeddings[-1] = image_embeddings[-1] + self.visual_model.no_mem_embed
|
|
|
|
feats = [
|
|
feat.permute(1, 2, 0).view(batch_size, -1, *feat_size)
|
|
for feat, feat_size in zip(image_embeddings[::-1], self._bb_feat_sizes[::-1])
|
|
][::-1]
|
|
_features = {"image_embed": feats[-1], "high_res_feats": feats[:-1]}
|
|
|
|
|
|
multimask_output = multimask_output
|
|
|
|
output = self.mm_extractor.beit3(visual_tokens=images_evf, textual_tokens=input_ids, text_padding_position=torch.zeros_like(input_ids))
|
|
|
|
feat = output["encoder_out"][:, :1, ...]
|
|
feat = self.text_hidden_fcs[0](feat)
|
|
(
|
|
sparse_embeddings,
|
|
dense_embeddings,
|
|
) = self.visual_model.sam_prompt_encoder(
|
|
points=None,
|
|
boxes=None,
|
|
masks=None,
|
|
text_embeds=feat,
|
|
)
|
|
high_res_features = _features["high_res_feats"]
|
|
sparse_embeddings = sparse_embeddings.to(feat.dtype)
|
|
low_res_masks, iou_predictions, _, _ = self.visual_model.sam_mask_decoder(
|
|
image_embeddings=_features["image_embed"],
|
|
image_pe=self.visual_model.sam_prompt_encoder.get_dense_pe(),
|
|
sparse_prompt_embeddings=sparse_embeddings,
|
|
dense_prompt_embeddings=dense_embeddings,
|
|
multimask_output=multimask_output,
|
|
repeat_image = True,
|
|
high_res_features=high_res_features,
|
|
)
|
|
if multimask_output:
|
|
sorted_ids = torch.argsort(iou_predictions, dim=-1, descending=True)
|
|
low_res_masks = torch.take_along_dim(low_res_masks, sorted_ids[..., None, None], dim=1)[:, :1]
|
|
|
|
pred_mask = self.postprocess_masks(
|
|
low_res_masks,
|
|
orig_hw=original_size_list[0],
|
|
)
|
|
|
|
return pred_mask[:, 0]
|
|
|
|
|
|
AutoConfig.register("evf", EvfConfig)
|
|
AutoModelForCausalLM.register(EvfConfig, EvfSam2Model) |