|
|
import os |
|
|
import os.path as osp |
|
|
import re |
|
|
|
|
|
from .base_eval_dataset import BaseEvalDataset |
|
|
|
|
|
from xtuner.registry import BUILDER |
|
|
import json |
|
|
from mmengine.dist import (master_only) |
|
|
from .textvqa_utils import TextVQAAccuracyEvaluator |
|
|
from mmengine.logging import print_log |
|
|
from .utils import custom_data_process |
|
|
|
|
|
def prompt_processor(prompt): |
|
|
if prompt.startswith('OCR tokens: '): |
|
|
pattern = r"Question: (.*?) Short answer:" |
|
|
match = re.search(pattern, prompt, re.DOTALL) |
|
|
question = match.group(1) |
|
|
elif 'Reference OCR token: ' in prompt and len(prompt.split('\n')) == 3: |
|
|
if prompt.startswith('Reference OCR token:'): |
|
|
question = prompt.split('\n')[1] |
|
|
else: |
|
|
question = prompt.split('\n')[0] |
|
|
elif len(prompt.split('\n')) == 2: |
|
|
question = prompt.split('\n')[0] |
|
|
else: |
|
|
assert False |
|
|
|
|
|
return question.lower() |
|
|
|
|
|
|
|
|
class TextVQADataset(BaseEvalDataset): |
|
|
METAINFO: dict = dict(name='textvqa') |
|
|
|
|
|
def __init__(self, data_file, ann_file, |
|
|
image_folder, image_processor, |
|
|
pad_image_to_square=True, metainfo=None,): |
|
|
super().__init__(metainfo) |
|
|
self.data_file = data_file |
|
|
self.ann_file = ann_file |
|
|
self.image_folder = image_folder |
|
|
|
|
|
self.image_processor = BUILDER.build(image_processor) |
|
|
self.pad_image_to_square = pad_image_to_square |
|
|
self.name = os.path.splitext(os.path.basename(data_file))[0] |
|
|
self.results_path = os.path.splitext(os.path.basename(data_file))[0] + '-results.jsonl' |
|
|
self.data = self.load_data_list() |
|
|
|
|
|
def load_data_list(self): |
|
|
data = [json.loads(q) for q in open(os.path.expanduser(self.data_file), "r")] |
|
|
for i, d in enumerate(data): |
|
|
d['img_id'] = i |
|
|
d['image_path'] = d['image'] |
|
|
d['question'] = d['text'] |
|
|
return data |
|
|
|
|
|
def __len__(self): |
|
|
return len(self.data) |
|
|
|
|
|
def __getitem__(self, idx): |
|
|
data = self.data[idx] |
|
|
data_dict = custom_data_process(self, data) |
|
|
return data_dict |
|
|
|
|
|
@master_only |
|
|
def evaluate(self, result, work_dir, show=True): |
|
|
answers_file = osp.join(work_dir, self.results_path) |
|
|
ans_file = open(answers_file, "w") |
|
|
|
|
|
for pred_dict in result: |
|
|
idx = pred_dict["img_id"] |
|
|
gt_data = self.data[idx] |
|
|
|
|
|
ans_file.write(json.dumps({"question_id": gt_data['question_id'], |
|
|
"prompt": gt_data['text'], |
|
|
"text": pred_dict['prediction'], |
|
|
"metadata": {}}) + "\n") |
|
|
ans_file.close() |
|
|
|
|
|
annotations = json.load(open(self.ann_file))['data'] |
|
|
annotations = {(annotation['image_id'], annotation['question'].lower()): annotation for annotation in |
|
|
annotations} |
|
|
results = [json.loads(line) for line in open(answers_file)] |
|
|
|
|
|
pred_list = [] |
|
|
for result in results: |
|
|
annotation = annotations[(result['question_id'], prompt_processor(result['prompt']))] |
|
|
pred_list.append({ |
|
|
"pred_answer": result['text'], |
|
|
"gt_answers": annotation['answers'], |
|
|
}) |
|
|
|
|
|
evaluator = TextVQAAccuracyEvaluator() |
|
|
acc = 100. * evaluator.eval_pred_list(pred_list) |
|
|
print_log('============================================', 'current') |
|
|
print_log('Samples: {}, Accuracy: {:.2f}%'.format(len(pred_list), acc), 'current') |
|
|
print_log('============================================', 'current') |
|
|
print_log(f'TextVQA successfully finished evaluating', 'current') |
|
|
return {'acc': acc} |
|
|
|
