|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.nn.functional as F |
|
|
from xtuner.registry import BUILDER |
|
|
from mmengine.model import BaseModel |
|
|
from xtuner.model.utils import guess_load_checkpoint |
|
|
from .utils import compute_mask_IoU |
|
|
|
|
|
|
|
|
class FrozenLlava(BaseModel): |
|
|
|
|
|
def __init__(self, |
|
|
model, |
|
|
mask_head, |
|
|
merge='mean', |
|
|
loss_mask=None, |
|
|
loss_dice=None, |
|
|
pretrained=None, |
|
|
**kwargs): |
|
|
super().__init__() |
|
|
self.llava = BUILDER.build(model) |
|
|
self.llava.requires_grad_(False) |
|
|
in_channels = (self.llava.config.text_config.num_attention_heads * |
|
|
self.llava.config.text_config.num_hidden_layers) |
|
|
mask_head.update(in_channels=in_channels) |
|
|
self.mask_head = BUILDER.build(mask_head) |
|
|
self.patch_size = self.llava.config.vision_config.patch_size |
|
|
self.merge = merge |
|
|
assert merge in ['mean', 'max'] |
|
|
|
|
|
self.loss_mask = BUILDER.build(loss_mask) |
|
|
self.loss_dice = BUILDER.build(loss_dice) |
|
|
|
|
|
self.text_layer_weights = nn.Parameter( |
|
|
torch.ones(self.llava.config.text_config.num_hidden_layers)) |
|
|
|
|
|
if pretrained is not None: |
|
|
_ = self.load_state_dict( |
|
|
guess_load_checkpoint(pretrained), strict=False) |
|
|
|
|
|
def get_text_layer_weights(self): |
|
|
return torch.softmax(self.text_layer_weights, dim=0) |
|
|
|
|
|
def apply_merge(self, x, dim=1): |
|
|
if self.merge == 'mean': |
|
|
return x.mean(dim=dim) |
|
|
elif self.merge == 'max': |
|
|
return x.max(dim=dim).values |
|
|
else: |
|
|
raise NotImplementedError |
|
|
|
|
|
def init_weights(self): |
|
|
pass |
|
|
|
|
|
def train(self, mode=True): |
|
|
super().train(mode=mode) |
|
|
self.llava.train(mode=False) |
|
|
self.training = mode |
|
|
return self |
|
|
|
|
|
def forward(self, data, data_samples=None, mode='loss'): |
|
|
if mode == 'loss': |
|
|
return self.compute_loss(data) |
|
|
elif mode == 'predict': |
|
|
return self.predict(data) |
|
|
elif mode == 'tensor': |
|
|
return self._forward(data) |
|
|
else: |
|
|
raise NotImplementedError |
|
|
|
|
|
def _compute(self, pred_masks, gt_masks): |
|
|
mask_cnt = pred_masks.shape[0] |
|
|
loss_dice = self.loss_dice( |
|
|
pred_masks.view(mask_cnt, -1), gt_masks.view(mask_cnt, -1), |
|
|
avg_factor=mask_cnt) |
|
|
loss_mask = self.loss_mask( |
|
|
pred_masks.view(-1), |
|
|
gt_masks.view(-1), |
|
|
avg_factor=pred_masks.numel()) |
|
|
accuracy = torch.eq((pred_masks.detach().sigmoid() > 0.5).to(gt_masks), |
|
|
gt_masks).to(gt_masks).mean() |
|
|
aiou = compute_mask_IoU((pred_masks.detach().sigmoid() > 0.5).to(gt_masks).view(mask_cnt, -1), |
|
|
gt_masks.view(mask_cnt, -1)).mean() |
|
|
|
|
|
return loss_dice, loss_mask, accuracy, aiou |
|
|
|
|
|
|
|
|
class FrozenLlavaSAM(FrozenLlava): |
|
|
def __init__(self, sam, *args, **kwargs): |
|
|
pretrained = kwargs.pop('pretrained', None) |
|
|
super().__init__(*args, **kwargs) |
|
|
self.sam = BUILDER.build(sam) |
|
|
self.text_proj = nn.Linear(self.llava.config.text_config.hidden_size, |
|
|
self.sam.model.prompt_encoder.embed_dim) |
|
|
|
|
|
if pretrained is not None: |
|
|
_ = self.load_state_dict( |
|
|
guess_load_checkpoint(pretrained), strict=False) |
|
|
|
|
|
def _forward(self, data_sample): |
|
|
text_layer_weights = self.get_text_layer_weights() |
|
|
inputs = dict(input_ids=data_sample['input_ids'][None].to(self.llava.device), |
|
|
mask_ids=data_sample['mask_ids'][None].to( |
|
|
self.llava.device), |
|
|
pixel_values=data_sample['pixel_values'][None].to(device=self.llava.device, |
|
|
dtype=self.llava.dtype), |
|
|
labels=data_sample['labels'][None].to(self.llava.device) |
|
|
) |
|
|
attention_mask = torch.ones(inputs['input_ids'].shape, device=self.llava.device, |
|
|
dtype=torch.bool) |
|
|
meta_data = data_sample['meta_data'] |
|
|
with torch.no_grad(): |
|
|
outputs = self.llava(**inputs, |
|
|
attention_mask=attention_mask, |
|
|
output_hidden_states=True, |
|
|
output_attentions=True) |
|
|
mask_ids = outputs['mask_ids'][0] |
|
|
attentions = [attn[0, ..., outputs['image_to_overwrite'][0]] |
|
|
for attn in outputs.attentions] |
|
|
hidden_states = outputs.hidden_states[-self.llava.config.text_config.num_hidden_layers:] |
|
|
|
|
|
labels = outputs.labels[0] |
|
|
|
|
|
|
|
|
hidden_states = torch.stack([hs[0] for hs in hidden_states]) |
|
|
|
|
|
hidden_states = ( |
|
|
hidden_states * text_layer_weights.view(-1, 1, 1)).sum(0) |
|
|
|
|
|
del outputs |
|
|
|
|
|
padded_h, padded_w = meta_data['padded_shape']['height'], meta_data['padded_shape']['width'] |
|
|
llava_h, llava_w = padded_h // self.patch_size, padded_w // self.patch_size |
|
|
|
|
|
attentions = [attn.view(*attn.shape[:-1], llava_h, llava_w) |
|
|
for attn in attentions] |
|
|
masks = data_sample['masks'] |
|
|
mask_attentions = [] |
|
|
text_embeds = [] |
|
|
for mask_id in range(len(masks)): |
|
|
matched = mask_ids == mask_id |
|
|
assert matched.sum() > 0 |
|
|
mask_attentions.append(torch.cat( |
|
|
[self.apply_merge(attn[:, matched], dim=1) for attn in attentions])) |
|
|
text_embeds.append(self.text_proj(hidden_states[matched])) |
|
|
|
|
|
del attentions |
|
|
mask_attentions = torch.stack(mask_attentions).to(self.mask_head.dtype) |
|
|
|
|
|
|
|
|
pred_masks = self.mask_head(mask_attentions)[:, 0] |
|
|
|
|
|
padded_mask_h, padded_mask_w = pred_masks.shape[-2:] |
|
|
|
|
|
before_height = int( |
|
|
meta_data['padding']['before_height'] * padded_mask_h / padded_h) |
|
|
before_width = int( |
|
|
meta_data['padding']['before_width'] * padded_mask_w / padded_w) |
|
|
|
|
|
mask_h = int(meta_data['image_shape']['height'] |
|
|
* padded_mask_h / padded_h + 0.5) |
|
|
mask_w = int(meta_data['image_shape']['width'] |
|
|
* padded_mask_w / padded_w + 0.5) |
|
|
pred_masks \ |
|
|
= pred_masks[:, before_height:before_height + mask_h, before_width:before_width + mask_w].contiguous() |
|
|
sam_pred_masks = self.sam( |
|
|
data_sample['image'], pred_masks, text_embeds) |
|
|
|
|
|
output = dict(pred_masks=pred_masks, sam_pred_masks=sam_pred_masks, |
|
|
labels=labels, mask_ids=mask_ids, hidden_states=hidden_states) |
|
|
|
|
|
return output |
|
|
|
|
|
@torch.no_grad() |
|
|
def predict(self, data_sample): |
|
|
return self._forward(data_sample)['sam_pred_masks'] |
|
|
|
|
|
def compute_loss(self, data): |
|
|
mask_cnts = 0 |
|
|
|
|
|
loss_dice = 0 |
|
|
loss_mask = 0 |
|
|
accuracy = 0 |
|
|
aiou = 0 |
|
|
|
|
|
sam_loss_dice = 0 |
|
|
sam_loss_mask = 0 |
|
|
sam_accuracy = 0 |
|
|
sam_aiou = 0 |
|
|
|
|
|
for data_sample in data: |
|
|
forward_output = self._forward(data_sample) |
|
|
pred_masks, sam_pred_masks = forward_output['pred_masks'], forward_output['sam_pred_masks'] |
|
|
masks = data_sample['masks'].to(self.llava.device) |
|
|
gt_masks = F.interpolate(masks[None].float(), |
|
|
size=pred_masks.shape[-2:])[0].to(pred_masks) |
|
|
sam_gt_masks = F.interpolate(masks[None].float(), |
|
|
size=sam_pred_masks.shape[-2:])[0].to(sam_pred_masks) |
|
|
|
|
|
mask_cnt = pred_masks.shape[0] |
|
|
assert pred_masks.shape == gt_masks.shape |
|
|
mask_cnts += mask_cnt |
|
|
|
|
|
loss_dice_, loss_mask_, accuracy_, aiou_ = self._compute( |
|
|
pred_masks, gt_masks) |
|
|
loss_dice += loss_dice_ * mask_cnt |
|
|
loss_mask += loss_mask_ * mask_cnt |
|
|
accuracy += accuracy_ * mask_cnt |
|
|
aiou += aiou_ * mask_cnt |
|
|
|
|
|
sam_loss_dice_, sam_loss_mask_, sam_accuracy_, sam_aiou_ = self._compute( |
|
|
sam_pred_masks, sam_gt_masks) |
|
|
sam_loss_dice += sam_loss_dice_ * mask_cnt |
|
|
sam_loss_mask += sam_loss_mask_ * mask_cnt |
|
|
sam_accuracy += sam_accuracy_ * mask_cnt |
|
|
sam_aiou += sam_aiou_ * mask_cnt |
|
|
|
|
|
assert mask_cnts > 0 |
|
|
|
|
|
loss_dict = {'loss_mask': loss_mask / mask_cnts, |
|
|
'loss_dice': loss_dice / mask_cnts, |
|
|
'accuracy': accuracy / mask_cnts, |
|
|
'aiou': aiou / mask_cnts, |
|
|
'sam_loss_mask': sam_loss_mask / mask_cnts, |
|
|
'sam_loss_dice': sam_loss_dice / mask_cnts, |
|
|
'sam_accuracy': sam_accuracy / mask_cnts, |
|
|
'sam_aiou': sam_aiou / mask_cnts, |
|
|
} |
|
|
|
|
|
return loss_dict |
|
|
|
