elia / test_elia.py
yxchng
add files
a166479
import datetime
import os
import time
import torch
import torch.utils.data
from torch import nn
from bert.multimodal_bert import MultiModalBert
import torchvision
from lib import multimodal_segmentation_ppm
import transforms as T
import utils
import numpy as np
from PIL import Image
import torch.nn.functional as F
from modeling.MaskFormerModel import MaskFormerHead
from addict import Dict
from bert.modeling_bert import BertLMPredictionHead, BertEncoder
def get_dataset(image_set, transform, args):
from data.dataset_refer_bert import ReferDataset
ds = ReferDataset(args,
split=image_set,
image_transforms=transform,
target_transforms=None,
eval_mode=True
)
num_classes = 2
return ds, num_classes
def evaluate(model, data_loader, device):
model.eval()
metric_logger = utils.MetricLogger(delimiter=" ")
# evaluation variables
cum_I, cum_U = 0, 0
eval_seg_iou_list = [.5, .6, .7, .8, .9]
seg_correct = np.zeros(len(eval_seg_iou_list), dtype=np.int32)
seg_total = 0
mean_IoU = []
header = 'Test:'
with torch.no_grad():
for data in metric_logger.log_every(data_loader, 100, header):
image, target, sentences, attentions = data
image, target, sentences, attentions = image.to(device), target.to(device), \
sentences.to(device), attentions.to(device)
sentences = sentences.squeeze(1)
attentions = attentions.squeeze(1)
target = target.cpu().data.numpy()
for j in range(sentences.size(-1)):
#if bert_model is not None:
# last_hidden_states = bert_model(sentences[:, :, j], attention_mask=attentions[:, :, j])[0]
# embedding = last_hidden_states.permute(0, 2, 1)
# output = model(image, embedding, l_mask=attentions[:, :, j].unsqueeze(-1))
#else:
output = model(image, sentences[:, :, j], attentions[:, :, j])
mask_cls_results = output["pred_logits"]
mask_pred_results = output["pred_masks"]
target_shape = target.shape[-2:]
mask_pred_results = F.interpolate(mask_pred_results, size=target_shape, mode='bilinear', align_corners=True)
pred_masks = model.semantic_inference(mask_cls_results, mask_pred_results)
output = pred_masks[0]
output = output.cpu()
#print(output.shape)
#output_mask = output.argmax(1).data.numpy()
output_mask = (output > 0.5).data.numpy()
I, U = computeIoU(output_mask, target)
if U == 0:
this_iou = 0.0
else:
this_iou = I*1.0/U
mean_IoU.append(this_iou)
cum_I += I
cum_U += U
for n_eval_iou in range(len(eval_seg_iou_list)):
eval_seg_iou = eval_seg_iou_list[n_eval_iou]
seg_correct[n_eval_iou] += (this_iou >= eval_seg_iou)
seg_total += 1
#del image, target, sentences, attentions, output, output_mask
#if bert_model is not None:
# del last_hidden_states, embedding
mean_IoU = np.array(mean_IoU)
mIoU = np.mean(mean_IoU)
print('Final results:')
print('Mean IoU is %.2f\n' % (mIoU*100.))
results_str = ''
for n_eval_iou in range(len(eval_seg_iou_list)):
results_str += ' precision@%s = %.2f\n' % \
(str(eval_seg_iou_list[n_eval_iou]), seg_correct[n_eval_iou] * 100. / seg_total)
results_str += ' overall IoU = %.2f\n' % (cum_I * 100. / cum_U)
print(results_str)
def get_transform(args):
transforms = [T.Resize(args.img_size, args.img_size),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
]
return T.Compose(transforms)
def computeIoU(pred_seg, gd_seg):
I = np.sum(np.logical_and(pred_seg, gd_seg))
U = np.sum(np.logical_or(pred_seg, gd_seg))
return I, U
class WrapperModel(nn.Module):
def __init__(self, image_model, language_model, classifier, args) :
super(WrapperModel, self).__init__()
self.image_model = image_model
self.language_model = language_model
self.classifier = classifier
self.lang_proj = nn.Linear(768,256)
config = Dict({
"architectures": [
"BertForMaskedLM"
],
"attention_probs_dropout_prob": 0.1,
"gradient_checkpointing": False,
"hidden_act": "gelu",
"hidden_dropout_prob": 0.1,
"hidden_size": 512,
"initializer_range": 0.02,
"intermediate_size": 3072,
"layer_norm_eps": 1e-12,
#"max_position_embeddings": 16+20,
"model_type": "bert",
"num_attention_heads": 8,
"num_hidden_layers": 8,
"pad_token_id": 0,
"position_embedding_type": "absolute",
"transformers_version": "4.6.0.dev0",
"type_vocab_size": 2,
"use_cache": True,
"vocab_size": 30522
})
self.mlm_transformer = BertEncoder(config)
self.lang_proj = nn.Linear(768,256)
self.mlm_vis_proj = nn.Conv2d(1024,512,1)
self.mlm_lang_proj = nn.Linear(768,512)
#print(vis_proj)
self.mlm_head = BertLMPredictionHead(config)
assert args.img_size % 4 == 0
num_img_tokens = 20 + ((args.img_size // 4)//8) ** 2
print(num_img_tokens)
self.mlm_pos_embeds = nn.Embedding(num_img_tokens+1, 512)
self.mlm_modal_embeds = nn.Embedding(3, 512)
self.mlm_mask_embed = nn.Embedding(1, 512)
self.mlm_pos_mlp = nn.Sequential(
nn.Linear(2, 512),
nn.LayerNorm(512),
nn.Linear(512,512),
nn.GELU()
)
def _get_binary_mask(self, target):
# 返回每类的binary mask
y, x = target.size()
target_onehot = torch.zeros(self.num_classes + 1, y, x)
target_onehot = target_onehot.scatter(dim=0, index=target.unsqueeze(0), value=1)
return target_onehot[1:]
def semantic_inference(self, mask_cls, mask_pred):
mask_cls = F.softmax(mask_cls, dim=1)[...,1:]
mask_pred = mask_pred.sigmoid()
semseg = torch.einsum("bqc,bqhw->bchw", mask_cls, mask_pred)
return semseg
def forward(self, image, sentences, attentions):
input_shape = image.shape[-2:]
l_mask = attentions.unsqueeze(dim=-1)
i0, Wh, Ww = self.image_model.forward_stem(image)
l0, extended_attention_mask = self.language_model.forward_stem(sentences, attentions)
i1 = self.image_model.forward_stage1(i0, Wh, Ww)
l1 = self.language_model.forward_stage1(l0, extended_attention_mask)
i1_residual, H, W, i1_temp, Wh, Ww = self.image_model.forward_pwam1(i1, Wh, Ww, l1, l_mask)
l1_residual, l1 = self.language_model.forward_pwam1(i1, l1, extended_attention_mask)
i1 = i1_temp
i2 = self.image_model.forward_stage2(i1, Wh, Ww)
l2 = self.language_model.forward_stage2(l1, extended_attention_mask)
i2_residual, H, W, i2_temp, Wh, Ww = self.image_model.forward_pwam2(i2, Wh, Ww, l2, l_mask)
l2_residual, l2 = self.language_model.forward_pwam2(i2, l2, extended_attention_mask)
i2 = i2_temp
i3 = self.image_model.forward_stage3(i2, Wh, Ww)
l3 = self.language_model.forward_stage3(l2, extended_attention_mask)
i3_residual, H, W, i3_temp, Wh, Ww = self.image_model.forward_pwam3(i3, Wh, Ww, l3, l_mask)
l3_residual, l3 = self.language_model.forward_pwam3(i3, l3, extended_attention_mask)
i3 = i3_temp
i4 = self.image_model.forward_stage4(i3, Wh, Ww)
l4 = self.language_model.forward_stage4(l3, extended_attention_mask)
i4_residual, H, W, i4_temp, Wh, Ww = self.image_model.forward_pwam4(i4, Wh, Ww, l4, l_mask)
l4_residual, l4 = self.language_model.forward_pwam4(i4, l4, extended_attention_mask)
i4 = i4_temp
#i1_residual, i2_residual, i3_residual, i4_residual = features
#x = self.classifier(i4_residual, i3_residual, i2_residual, i1_residual)
#x = F.interpolate(x, size=input_shape, mode='bilinear', align_corners=True)
outputs = {}
outputs['s1'] = i1_residual
outputs['s2'] = i2_residual
outputs['s3'] = i3_residual
outputs['s4'] = i4_residual
predictions, _ = self.classifier(outputs)
return predictions
def main(args):
#def main(local_rank, args):
#device = torch.device(args.device)
device = 'cuda'
dataset_test, _ = get_dataset(args.split, get_transform(args=args), args)
test_sampler = torch.utils.data.SequentialSampler(dataset_test)
data_loader_test = torch.utils.data.DataLoader(dataset_test, batch_size=1,
sampler=test_sampler, num_workers=args.workers)
print(args.model)
single_model = multimodal_segmentation_ppm.__dict__[args.model](pretrained='',args=args)
#single_model = MultiModalFocal(depths=[2, 2, 18, 2], embed_dim=128, focal_levels=[3, 3, 3, 3], focal_windows=[9,9,9,9], drop_path_rate=0.3)
#single_model.init_weights('./focalnet_base_lrf.pth')
checkpoint = torch.load(args.resume, map_location='cpu')
#single_model.load_state_dict(checkpoint['model'])
#model = single_model.to(device)
if args.model != 'lavt_one':
model_class = MultiModalBert
#single_bert_model = model_class.from_pretrained(args.ck_bert, embed_dim=128)
single_bert_model = model_class.from_pretrained(args.ck_bert, embed_dim=single_model.backbone.embed_dim)
# work-around for a transformers bug; need to update to a newer version of transformers to remove these two lines
if args.ddp_trained_weights:
single_bert_model.pooler = None
#single_bert_model.load_state_dict(checkpoint['bert_model'])
#bert_model = single_bert_model.to(device)
else:
bert_model = None
#model = WrapperModel(single_model.backbone, single_bert_model, single_model.classifier)
#model.load_state_dict(checkpoint['model'])
#model.to(device)
input_shape = dict()
input_shape['s1'] = Dict({'channel': 128, 'stride': 4})
input_shape['s2'] = Dict({'channel': 256, 'stride': 8})
input_shape['s3'] = Dict({'channel': 512, 'stride': 16})
input_shape['s4'] = Dict({'channel': 1024, 'stride': 32})
cfg = Dict()
cfg.MODEL.SEM_SEG_HEAD.COMMON_STRIDE = 4
cfg.MODEL.MASK_FORMER.DROPOUT = 0.0
cfg.MODEL.MASK_FORMER.NHEADS = 8
cfg.MODEL.SEM_SEG_HEAD.TRANSFORMER_ENC_LAYERS = 4
cfg.MODEL.SEM_SEG_HEAD.CONVS_DIM = 256
cfg.MODEL.SEM_SEG_HEAD.MASK_DIM = 256
cfg.MODEL.SEM_SEG_HEAD.DEFORMABLE_TRANSFORMER_ENCODER_IN_FEATURES = ["s1", "s2", "s3", "s4"]
cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES = 1
cfg.MODEL.MASK_FORMER.HIDDEN_DIM = 256
cfg.MODEL.MASK_FORMER.NUM_OBJECT_QUERIES = 1
cfg.MODEL.MASK_FORMER.DIM_FEEDFORWARD = 2048
cfg.MODEL.MASK_FORMER.DEC_LAYERS = 10
cfg.MODEL.MASK_FORMER.PRE_NORM = False
maskformer_head = MaskFormerHead(cfg, input_shape)
#maskformer_head = torch.nn.SyncBatchNorm.convert_sync_batchnorm(maskformer_head)
#maskformer_head.cuda()
#maskformer_head = torch.nn.parallel.DistributedDataParallel(maskformer_head, device_ids=[args.local_rank], find_unused_parameters=False)
#single_head = maskformer_head.module
#print(single_head)
model = WrapperModel(single_model.backbone, single_bert_model, maskformer_head, args)
model.load_state_dict(checkpoint['model'])
model.to(device)
#model.cuda()
#model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], find_unused_parameters=True)
#single_model = model.module
#model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], find_unused_parameters=True)
#single_model = model.module
evaluate(model, data_loader_test, device=device)
if __name__ == "__main__":
from args import get_parser
parser = get_parser()
args = parser.parse_args()
print('Image size: {}'.format(str(args.img_size)))
print(args)
main(args)
#mp.spawn(main, args=(args,), nprocs=torch.cuda.device_count())