| import json | |
| import os | |
| import re | |
| from torch.utils.data import Dataset | |
| 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(Dataset): | |
| def __init__( | |
| self, | |
| image_dir="./downloads/TextVQA/train_images", | |
| ann_path="./downloads/TextVQA/TextVQA_0.5.1_val.json", | |
| ): | |
| self.data = json.load(open(ann_path, "r"))["data"] | |
| self.image_dir = image_dir | |
| def __len__(self): | |
| return len(self.data) | |
| def __getitem__(self, idx): | |
| question = self.data[idx]['question'] | |
| answers = self.data[idx]['answers'] | |
| img_id = self.data[idx]['image_id'] | |
| qid = self.data[idx]['question_id'] | |
| img_path = os.path.join(self.image_dir, f"{img_id}.jpg") | |
| item = { | |
| "question_id": qid, | |
| "image_path": img_path, | |
| "question": question, | |
| "gt_answers": answers | |
| } | |
| return item | |
| class docVQADataset(Dataset): | |
| def __init__( | |
| self, | |
| image_dir= "./downloads/DocVQA/spdocvqa_images", | |
| ann_path= "./downloads/DocVQA/val_v1.0_withQT.json", | |
| ocr_token_path=None | |
| ): | |
| self.data = json.load(open(ann_path, "r"))["data"] | |
| self.image_dir = image_dir | |
| self.ann_path = ann_path | |
| if ocr_token_path: | |
| self.ocr_token_data = {item['image_id']: item for item in json.load(open(ocr_token_path, "r"))["data"]} | |
| def __len__(self): | |
| return len(self.data) | |
| def __getitem__(self, idx): | |
| question_id = self.data[idx]['questionId'] | |
| relative_img_path = self.data[idx]['image'] | |
| corrected_relative_img_path = relative_img_path.replace("documents", "images") | |
| img_path = os.path.join(self.image_dir, corrected_relative_img_path) | |
| question = self.data[idx]['question'] | |
| answers = self.data[idx]['answers'] | |
| question_type = self.data[idx]['question_types'] | |
| return { | |
| "question_id": question_id, | |
| "image_path": img_path, | |
| "question": question, | |
| "gt_answers": answers, | |
| 'question_type': question_type, | |
| } | |
| class docVQATESTDataset(Dataset): | |
| def __init__( | |
| self, | |
| image_dir= "./downloads/DocVQA/spdocvqa_images", | |
| ann_path= "./downloads/DocVQA/test_v1.0.json", | |
| ocr_token_path=None | |
| ): | |
| self.data = json.load(open(ann_path, "r"))["data"] | |
| self.image_dir = image_dir | |
| self.ann_path = ann_path | |
| def __len__(self): | |
| return len(self.data) | |
| def __getitem__(self, idx): | |
| question_id = self.data[idx]['questionId'] | |
| relative_img_path = self.data[idx]['image'] | |
| corrected_relative_img_path = relative_img_path.replace("documents", "images") | |
| img_path = os.path.join(self.image_dir, corrected_relative_img_path) | |
| question = self.data[idx]['question'] | |
| return { | |
| "question_id": question_id, | |
| "image_path": img_path, | |
| "question": question, | |
| "gt_answers": "", | |
| 'question_type': "", | |
| } | |