395 lines
14 KiB
Python
395 lines
14 KiB
Python
import argparse
|
|
import itertools
|
|
import json
|
|
import os
|
|
import random
|
|
import time
|
|
from functools import partial
|
|
from typing import Optional
|
|
|
|
import torch
|
|
from tqdm import tqdm
|
|
from PIL import Image
|
|
|
|
from mplug_owl2.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
|
|
from mplug_owl2.conversation import conv_templates, SeparatorStyle
|
|
from mplug_owl2.model.builder import load_pretrained_model
|
|
from mplug_owl2.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
|
|
|
|
from sklearn.metrics import accuracy_score, precision_score, recall_score, confusion_matrix
|
|
|
|
ds_collections = {
|
|
'mme': {
|
|
"test": "MME_Benchmark_release_version/eval_tool/Your_Results",
|
|
"base_dir": 'MME_Benchmark_release_version',
|
|
'max_new_tokens': 10,
|
|
},
|
|
}
|
|
|
|
|
|
eval_type_dict = {
|
|
"Perception": ["existence", "count", "position", "color", "posters", "celebrity", "scene", "landmark", "artwork", "OCR"],
|
|
"Cognition": ["commonsense_reasoning", "numerical_calculation", "text_translation", "code_reasoning"]
|
|
}
|
|
|
|
|
|
class calculate_metrics:
|
|
def divide_chunks(self, l, n=2):
|
|
# looping till length l
|
|
for i in range(0, len(l), n):
|
|
yield l[i:i + n]
|
|
|
|
return
|
|
|
|
def parse_pred_ans(self, pred_ans):
|
|
pred_label = None
|
|
if pred_ans in ["yes", "no"]:
|
|
pred_label = pred_ans
|
|
else:
|
|
prefix_pred_ans = pred_ans[:4]
|
|
|
|
if "yes" in prefix_pred_ans:
|
|
pred_label = "yes"
|
|
elif "no" in prefix_pred_ans:
|
|
pred_label = "no"
|
|
else:
|
|
pred_label = "other"
|
|
|
|
return pred_label
|
|
|
|
|
|
def compute_metric(self, gts, preds):
|
|
assert len(gts) == len(preds)
|
|
|
|
label_map = {
|
|
"yes": 1,
|
|
"no": 0,
|
|
"other": -1,
|
|
}
|
|
|
|
gts = [label_map[x] for x in gts]
|
|
preds = [label_map[x] for x in preds]
|
|
|
|
acc = accuracy_score(gts, preds)
|
|
|
|
clean_gts = []
|
|
clean_preds = []
|
|
other_num = 0
|
|
for gt, pred in zip(gts, preds):
|
|
if pred == -1:
|
|
other_num += 1
|
|
continue
|
|
clean_gts.append(gt)
|
|
clean_preds.append(pred)
|
|
|
|
|
|
conf_mat = confusion_matrix(clean_gts, clean_preds, labels=[1,0])
|
|
precision = precision_score(clean_gts, clean_preds, average='binary')
|
|
recall = recall_score(clean_gts, clean_preds, average='binary')
|
|
tp, fn = conf_mat[0]
|
|
fp, tn = conf_mat[1]
|
|
|
|
metric_dict = dict()
|
|
metric_dict = {
|
|
"TP": tp,
|
|
"FN": fn,
|
|
"TN": tn,
|
|
"FP": fp,
|
|
"precision": precision,
|
|
"recall": recall,
|
|
"other_num": other_num,
|
|
"acc": acc,
|
|
}
|
|
|
|
return metric_dict
|
|
|
|
|
|
def process_result(self, outputs):
|
|
|
|
model_score_dict = dict()
|
|
for eval_type, task_name_list in eval_type_dict.items():
|
|
print("===========", eval_type, "===========")
|
|
scores = 0
|
|
task_score_dict = dict()
|
|
|
|
for task_name in task_name_list:
|
|
chunk_lines = outputs[task_name]
|
|
img_num = len(chunk_lines)
|
|
task_other_ans_num = 0
|
|
task_score = 0
|
|
acc_plus_correct_num = 0
|
|
gts = []
|
|
preds = []
|
|
|
|
for k, img_items in chunk_lines.items():
|
|
assert len(img_items) == 2
|
|
img_correct_num = 0
|
|
|
|
for img_item in img_items:
|
|
img_name, question, gt_ans, pred_ans = img_item['image_name'], img_item['question'], img_item['annotation'], img_item['answer']
|
|
gt_ans = gt_ans.lower()
|
|
pred_ans = pred_ans.lower()
|
|
|
|
assert gt_ans in ["yes", "no"] # gt can only be yes or no.
|
|
|
|
pred_ans = self.parse_pred_ans(pred_ans)
|
|
assert pred_ans in ["yes", "no", "other"]
|
|
|
|
gts.append(gt_ans)
|
|
preds.append(pred_ans)
|
|
|
|
if gt_ans == pred_ans:
|
|
img_correct_num += 1
|
|
|
|
if pred_ans not in ["yes", "no"]:
|
|
task_other_ans_num += 1
|
|
|
|
if img_correct_num == 2:
|
|
acc_plus_correct_num += 1
|
|
|
|
# cal TP precision acc, etc.
|
|
metric_dict = self.compute_metric(gts, preds)
|
|
acc_plus = acc_plus_correct_num / img_num
|
|
metric_dict["acc_plus"] = acc_plus
|
|
|
|
|
|
for k, v in metric_dict.items():
|
|
if k in ["acc", "acc_plus"]:
|
|
task_score += v*100
|
|
|
|
task_score_dict[task_name] = task_score
|
|
|
|
scores += task_score
|
|
|
|
print("total score:", scores, "\n")
|
|
for task_name, score in task_score_dict.items():
|
|
print("\t", task_name, " score:", score)
|
|
model_score_dict[eval_type+"-"+task_name] = score
|
|
print("\n")
|
|
|
|
return model_score_dict
|
|
|
|
def collate_fn(batches, tokenizer):
|
|
|
|
questions = [_['question'] for _ in batches]
|
|
questions_origin = [_['question_origin'] for _ in batches]
|
|
annotations = [_['annotation'] for _ in batches]
|
|
image_names = [_['image_name'] for _ in batches]
|
|
categories = [_['category'] for _ in batches]
|
|
|
|
image_tensor = [_['image_tensor'] for _ in batches]
|
|
|
|
input_ids = []
|
|
for input_text in questions:
|
|
input_ids.append(tokenizer_image_token(input_text, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').tolist())
|
|
input_tokens_max_length = max([len(x) for x in input_ids])
|
|
pad_token_id = tokenizer.pad_token_id
|
|
|
|
input_ids = [([pad_token_id] * (input_tokens_max_length - len(_)) + _) for _ in input_ids] # pad in the left
|
|
input_ids = torch.LongTensor(input_ids)
|
|
attention_mask = 1 - input_ids.eq(pad_token_id).long()
|
|
|
|
image_tensor = torch.cat(image_tensor, dim=0)
|
|
return questions_origin, image_tensor, input_ids, attention_mask, annotations, categories, image_names
|
|
|
|
|
|
class VQADataset(torch.utils.data.Dataset):
|
|
|
|
def __init__(self, test, base_dir, prompt, image_processor):
|
|
|
|
self.prompt = prompt
|
|
self.image_processor = image_processor
|
|
|
|
self.categories = [
|
|
'OCR', 'artwork', 'celebrity',
|
|
'code_reasoning', 'color',
|
|
'commonsense_reasoning','count',
|
|
'existence', 'landmark',
|
|
'numerical_calculation', 'position',
|
|
'posters', 'scene', 'text_translation'
|
|
]
|
|
self.data = []
|
|
self.base_dir = base_dir
|
|
for category in self.categories:
|
|
anno_path = os.path.join(test, '{}.txt'.format(category))
|
|
with open(anno_path, 'r') as f:
|
|
self.data += [(x, category) for x in f.read().splitlines()]
|
|
|
|
def __len__(self):
|
|
return len(self.data)
|
|
|
|
def __getitem__(self, idx):
|
|
data = self.data[idx]
|
|
d, category = data
|
|
image_name, question, annotation = d.split('\t')
|
|
|
|
category_path = os.path.join(self.base_dir, category)
|
|
if category in ['landmark', 'artwork', 'celebrity', 'posters', 'scene']:
|
|
image_path = os.path.join(category_path, 'images')
|
|
text_path = os.path.join(category_path, 'questions_answers_YN')
|
|
else:
|
|
image_path = category_path
|
|
text_path = category_path
|
|
image = os.path.join(image_path, image_name)
|
|
|
|
image = Image.open(image).convert('RGB')
|
|
max_edge = max(image.size)
|
|
image = image.resize((max_edge, max_edge)) # Resize here for best performance
|
|
image_tensor = process_images([image], self.image_processor)
|
|
|
|
return {
|
|
'image_tensor': image_tensor,
|
|
'question': self.prompt.format(question).replace('Please answer yes or no.', ''),
|
|
"category": category,
|
|
"image_name": image_name,
|
|
'question_origin': question,
|
|
'annotation': annotation
|
|
}
|
|
|
|
|
|
class InferenceSampler(torch.utils.data.sampler.Sampler):
|
|
|
|
def __init__(self, size):
|
|
self._size = int(size)
|
|
assert size > 0
|
|
self._rank = torch.distributed.get_rank()
|
|
self._world_size = torch.distributed.get_world_size()
|
|
self._local_indices = self._get_local_indices(size, self._world_size,
|
|
self._rank)
|
|
|
|
@staticmethod
|
|
def _get_local_indices(total_size, world_size, rank):
|
|
shard_size = total_size // world_size
|
|
left = total_size % world_size
|
|
shard_sizes = [shard_size + int(r < left) for r in range(world_size)]
|
|
|
|
begin = sum(shard_sizes[:rank])
|
|
end = min(sum(shard_sizes[:rank + 1]), total_size)
|
|
return range(begin, end)
|
|
|
|
def __iter__(self):
|
|
yield from self._local_indices
|
|
|
|
def __len__(self):
|
|
return len(self._local_indices)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--checkpoint', type=str, default='')
|
|
parser.add_argument('--dataset', type=str, default='mme')
|
|
parser.add_argument('--batch-size', type=int, default=1)
|
|
parser.add_argument('--num-workers', type=int, default=1)
|
|
parser.add_argument('--seed', type=int, default=0)
|
|
args = parser.parse_args()
|
|
|
|
torch.distributed.init_process_group(
|
|
backend='nccl',
|
|
world_size=int(os.getenv('WORLD_SIZE', '1')),
|
|
rank=int(os.getenv('RANK', '0')),
|
|
)
|
|
|
|
torch.cuda.set_device(int(os.getenv('LOCAL_RANK', 0)))
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = os.getenv('LOCAL_RANK', "0")
|
|
|
|
model_path = args.checkpoint
|
|
model_name = get_model_name_from_path(model_path)
|
|
|
|
tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name, load_8bit=False, load_4bit=False, device_map={"":f"cuda:{os.getenv('LOCAL_RANK', '0')}"}, device="cuda")
|
|
tokenizer.padding_side = 'left'
|
|
if not hasattr(tokenizer, 'pad_token_id'):
|
|
tokenizer.pad_token_id = tokenizer.eos_token_id
|
|
|
|
prompt = 'USER: <|image|>{}\nAnswer the question using a single word or phrase. ASSISTANT: '
|
|
|
|
random.seed(args.seed)
|
|
dataset = VQADataset(
|
|
test=ds_collections[args.dataset]['test'],
|
|
base_dir=ds_collections[args.dataset]['base_dir'],
|
|
prompt=prompt,
|
|
image_processor=image_processor,
|
|
)
|
|
|
|
dataloader = torch.utils.data.DataLoader(
|
|
dataset=dataset,
|
|
sampler=InferenceSampler(len(dataset)),
|
|
batch_size=args.batch_size,
|
|
num_workers=args.num_workers,
|
|
pin_memory=True,
|
|
drop_last=False,
|
|
collate_fn=partial(collate_fn, tokenizer=tokenizer),
|
|
)
|
|
|
|
outputs = []
|
|
for _, (questions, image_tensor, input_ids, attention_mask,
|
|
annotations, categories, image_names) in tqdm(enumerate(dataloader)):
|
|
pred = model.generate(
|
|
input_ids=input_ids.cuda(),
|
|
attention_mask=attention_mask.cuda(),
|
|
images=image_tensor.to(dtype=model.dtype).cuda(),
|
|
do_sample=False,
|
|
num_beams=1,
|
|
max_new_tokens=ds_collections[args.dataset]['max_new_tokens'],
|
|
min_new_tokens=1,
|
|
length_penalty=1,
|
|
num_return_sequences=1,
|
|
output_hidden_states=True,
|
|
use_cache=True,
|
|
)
|
|
answers = [
|
|
tokenizer.decode(_[input_ids.size(1):].cpu(),
|
|
skip_special_tokens=True).strip() for _ in pred
|
|
]
|
|
|
|
for question, answer, annotation, category, image_name in zip(questions, answers,
|
|
annotations, categories, image_names):
|
|
outputs.append({
|
|
'image_name': image_name,
|
|
'category': category,
|
|
'question': question,
|
|
'answer': answer,
|
|
'annotation': annotation,
|
|
})
|
|
|
|
torch.distributed.barrier()
|
|
|
|
world_size = torch.distributed.get_world_size()
|
|
merged_outputs = [None for _ in range(world_size)]
|
|
torch.distributed.all_gather_object(merged_outputs, json.dumps(outputs))
|
|
|
|
merged_outputs = [json.loads(_) for _ in merged_outputs]
|
|
merged_outputs = [_ for _ in itertools.chain.from_iterable(merged_outputs)]
|
|
|
|
if torch.distributed.get_rank() == 0:
|
|
print(f"Evaluating {args.dataset} ...")
|
|
time_prefix = time.strftime('%y%m%d%H%M%S', time.localtime())
|
|
results_file = f'{args.dataset}_{time_prefix}_s{args.seed}.json'
|
|
json.dump(merged_outputs, open(results_file, 'w'), ensure_ascii=False)
|
|
|
|
|
|
groups = {}
|
|
for output in merged_outputs:
|
|
if not output['category'] in groups:
|
|
groups[output['category']] = {}
|
|
if not output['image_name'] in groups[output['category']]:
|
|
groups[output['category']][output['image_name']] = []
|
|
|
|
groups[output['category']][output['image_name']].append({
|
|
"image_name": output['image_name'],
|
|
"question": output['question'],
|
|
"answer": output['answer'],
|
|
"annotation": output['annotation'],
|
|
})
|
|
|
|
cal = calculate_metrics()
|
|
model_score_dict = cal.process_result(groups)
|
|
metrics = {
|
|
"Perception": sum([v for k, v in model_score_dict.items() if 'Perception' in k]),
|
|
"Cognition": sum([v for k, v in model_score_dict.items() if 'Cognition' in k]),
|
|
**model_score_dict
|
|
}
|
|
print(metrics)
|
|
|
|
torch.distributed.barrier()
|