| 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': "", |
| } |
|
|
