| import json |
| import os |
|
|
| from mmengine.dist import (master_only) |
|
|
| from vlm.datasets.evaluation.base_eval_dataset import BaseEvalDataset |
|
|
| from mmengine.logging import print_log |
|
|
|
|
| class FigureReasoningDataset(BaseEvalDataset): |
| METAINFO: dict = dict(name='FigureReasoning dataset') |
|
|
| def __init__(self, |
| data_file, |
| image_folder, |
| metainfo={}, |
| ): |
| super().__init__(metainfo) |
| self.data_file = data_file |
| with open(data_file, 'r') as f: |
| self.data = json.load(f) |
| self.image_folder = image_folder |
|
|
| self.type_map = { |
| 'yangshiguilv': 'Style', |
| 'shuliangguilv': 'Quantity', |
| 'weizhiguilv': 'Positional', |
| 'shuxingguilv': 'Attribute', |
| 'kongjianguilv': 'Spatial', |
| 'other': 'Others', |
| } |
| self.difficulty_nums = {'hard': 0, 'medium': 0, 'easy': 0} |
| self.difficulty_tp = {'hard': 0, 'medium': 0, 'easy': 0} |
| self.type_nums = {'Style': 0, 'Quantity': 0, 'Positional': 0, 'Attribute': 0, 'Spatial': 0, 'Others': 0} |
| self.type_tp = {'Style': 0, 'Quantity': 0, 'Positional': 0, 'Attribute': 0, 'Spatial': 0, 'Others': 0} |
| self.all_nums = 0 |
| self.all_tp = 0 |
|
|
| def __len__(self): |
| return len(self.data) |
|
|
| def __getitem__(self, idx): |
| data = self.data[idx] |
| data_type = data['type'] |
| difficulty = data['difficulty'] |
|
|
| data_type = data_type.split('/')[0] |
| data_type = self.type_map[data_type] |
|
|
| image = data['image'] |
| image_path = os.path.join(self.image_folder, image) |
|
|
| |
| question = data['conversations'][0]['value'] |
|
|
| answer = data['answer'] |
| ret = {'difficulty': difficulty, 'type': data_type, 'image_path': image_path, |
| 'question': question, 'answer': answer, |
| 'text_prompts': None, 'pixel_values': None, |
| 'img_id': data['id'], |
| } |
| return ret |
|
|
| def _get_answer(self, prediction): |
| answer_list = ['A', 'B', 'C', 'D'] |
| prediction = prediction.replace('.', ' ').split(' ') |
| for _aws in answer_list: |
| if _aws in prediction: |
| return _aws |
| return 'None' |
|
|
| @master_only |
| def evaluate(self, results, work_dir): |
| for result in results: |
| prediction = result['prediction'] |
| prediction = self._get_answer(prediction) |
|
|
| answer = result['answer'] |
| difficulty = result['difficulty'] |
| data_type = result['type'] |
|
|
| self.difficulty_nums[difficulty] += 1 |
| self.type_nums[data_type] += 1 |
| self.all_nums += 1 |
| |
| if prediction == answer: |
| |
| self.difficulty_tp[difficulty] += 1 |
| self.type_tp[data_type] += 1 |
| self.all_tp += 1 |
| print_log('============================================', 'current') |
| print_log(f'Figure reasoning dataset successfully finished evaluating, {len(results)} items' 'current') |
| print_log(f'Average: {self.all_tp / (self.all_nums + 1e-4)}', 'current') |
| for _type in self.type_nums.keys(): |
| print_log(f'{_type}: {self.type_tp[_type] / (self.type_nums[_type] + 1e-4)}', 'current') |
| for _type in self.difficulty_nums.keys(): |
| print_log(f'{_type}: {self.difficulty_tp[_type] / (self.difficulty_nums[_type] + 1e-4)}', 'current') |
| print_log('============================================', 'current') |
| return None |
|
|
