Spaces:
Running on Zero
Running on Zero
| import os | |
| import json | |
| import torch | |
| from torch.utils.data.dataset import Dataset | |
| from torchvision import transforms | |
| from PIL import Image | |
| from .image_utils import Compose, RandomResizedCropAndInterpolationWithTwoPic | |
| XFund_label2ids = { | |
| "O":0, | |
| 'B-HEADER':1, | |
| 'I-HEADER':2, | |
| 'B-QUESTION':3, | |
| 'I-QUESTION':4, | |
| 'B-ANSWER':5, | |
| 'I-ANSWER':6, | |
| } | |
| class xfund_dataset(Dataset): | |
| def box_norm(self, box, width, height): | |
| def clip(min_num, num, max_num): | |
| return min(max(num, min_num), max_num) | |
| x0, y0, x1, y1 = box | |
| x0 = clip(0, int((x0 / width) * 1000), 1000) | |
| y0 = clip(0, int((y0 / height) * 1000), 1000) | |
| x1 = clip(0, int((x1 / width) * 1000), 1000) | |
| y1 = clip(0, int((y1 / height) * 1000), 1000) | |
| assert x1 >= x0 | |
| assert y1 >= y0 | |
| return [x0, y0, x1, y1] | |
| def get_segment_ids(self, bboxs): | |
| segment_ids = [] | |
| for i in range(len(bboxs)): | |
| if i == 0: | |
| segment_ids.append(0) | |
| else: | |
| if bboxs[i - 1] == bboxs[i]: | |
| segment_ids.append(segment_ids[-1]) | |
| else: | |
| segment_ids.append(segment_ids[-1] + 1) | |
| return segment_ids | |
| def get_position_ids(self, segment_ids): | |
| position_ids = [] | |
| for i in range(len(segment_ids)): | |
| if i == 0: | |
| position_ids.append(2) | |
| else: | |
| if segment_ids[i] == segment_ids[i - 1]: | |
| position_ids.append(position_ids[-1] + 1) | |
| else: | |
| position_ids.append(2) | |
| return position_ids | |
| def load_data( | |
| self, | |
| data_file, | |
| ): | |
| # re-org data format | |
| total_data = {"id": [], "lines": [], "bboxes": [], "ner_tags": [], "image_path": []} | |
| for i in range(len(data_file['documents'])): | |
| width, height = data_file['documents'][i]['img']['width'], data_file['documents'][i]['img'][ | |
| 'height'] | |
| cur_doc_lines, cur_doc_bboxes, cur_doc_ner_tags, cur_doc_image_path = [], [], [], [] | |
| for j in range(len(data_file['documents'][i]['document'])): | |
| cur_item = data_file['documents'][i]['document'][j] | |
| cur_doc_lines.append(cur_item['text']) | |
| cur_doc_bboxes.append(self.box_norm(cur_item['box'], width=width, height=height)) | |
| cur_doc_ner_tags.append(cur_item['label']) | |
| total_data['id'] += [len(total_data['id'])] | |
| total_data['lines'] += [cur_doc_lines] | |
| total_data['bboxes'] += [cur_doc_bboxes] | |
| total_data['ner_tags'] += [cur_doc_ner_tags] | |
| total_data['image_path'] += [data_file['documents'][i]['img']['fname']] | |
| # tokenize text and get bbox/label | |
| total_input_ids, total_bboxs, total_label_ids = [], [], [] | |
| for i in range(len(total_data['lines'])): | |
| cur_doc_input_ids, cur_doc_bboxs, cur_doc_labels = [], [], [] | |
| for j in range(len(total_data['lines'][i])): | |
| cur_input_ids = self.tokenizer(total_data['lines'][i][j], truncation=False, add_special_tokens=False, return_attention_mask=False)['input_ids'] | |
| if len(cur_input_ids) == 0: continue | |
| cur_label = total_data['ner_tags'][i][j].upper() | |
| if cur_label == 'OTHER': | |
| cur_labels = ["O"] * len(cur_input_ids) | |
| for k in range(len(cur_labels)): | |
| cur_labels[k] = self.label2ids[cur_labels[k]] | |
| else: | |
| cur_labels = [cur_label] * len(cur_input_ids) | |
| cur_labels[0] = self.label2ids['B-' + cur_labels[0]] | |
| for k in range(1, len(cur_labels)): | |
| cur_labels[k] = self.label2ids['I-' + cur_labels[k]] | |
| assert len(cur_input_ids) == len([total_data['bboxes'][i][j]] * len(cur_input_ids)) == len(cur_labels) | |
| cur_doc_input_ids += cur_input_ids | |
| cur_doc_bboxs += [total_data['bboxes'][i][j]] * len(cur_input_ids) | |
| cur_doc_labels += cur_labels | |
| assert len(cur_doc_input_ids) == len(cur_doc_bboxs) == len(cur_doc_labels) | |
| assert len(cur_doc_input_ids) > 0 | |
| total_input_ids.append(cur_doc_input_ids) | |
| total_bboxs.append(cur_doc_bboxs) | |
| total_label_ids.append(cur_doc_labels) | |
| assert len(total_input_ids) == len(total_bboxs) == len(total_label_ids) | |
| # split text to several slices because of over-length | |
| input_ids, bboxs, labels = [], [], [] | |
| segment_ids, position_ids = [], [] | |
| image_path = [] | |
| for i in range(len(total_input_ids)): | |
| start = 0 | |
| cur_iter = 0 | |
| while start < len(total_input_ids[i]): | |
| end = min(start + 510, len(total_input_ids[i])) | |
| input_ids.append([self.tokenizer.cls_token_id] + total_input_ids[i][start: end] + [self.tokenizer.sep_token_id]) | |
| bboxs.append([[0, 0, 0, 0]] + total_bboxs[i][start: end] + [[1000, 1000, 1000, 1000]]) | |
| labels.append([-100] + total_label_ids[i][start: end] + [-100]) | |
| cur_segment_ids = self.get_segment_ids(bboxs[-1]) | |
| cur_position_ids = self.get_position_ids(cur_segment_ids) | |
| segment_ids.append(cur_segment_ids) | |
| position_ids.append(cur_position_ids) | |
| image_path.append(os.path.join(self.args.data_dir, "images", total_data['image_path'][i])) | |
| start = end | |
| cur_iter += 1 | |
| assert len(input_ids) == len(bboxs) == len(labels) == len(segment_ids) == len(position_ids) | |
| assert len(segment_ids) == len(image_path) | |
| res = { | |
| 'input_ids': input_ids, | |
| 'bbox': bboxs, | |
| 'labels': labels, | |
| 'segment_ids': segment_ids, | |
| 'position_ids': position_ids, | |
| 'image_path': image_path, | |
| } | |
| return res | |
| def __init__( | |
| self, | |
| args, | |
| tokenizer, | |
| mode | |
| ): | |
| self.args = args | |
| self.mode = mode | |
| self.cur_la = args.language | |
| self.tokenizer = tokenizer | |
| self.label2ids = XFund_label2ids | |
| self.common_transform = Compose([ | |
| RandomResizedCropAndInterpolationWithTwoPic( | |
| size=args.input_size, interpolation=args.train_interpolation, | |
| ), | |
| ]) | |
| self.patch_transform = transforms.Compose([ | |
| transforms.ToTensor(), | |
| transforms.Normalize( | |
| mean=torch.tensor((0.5, 0.5, 0.5)), | |
| std=torch.tensor((0.5, 0.5, 0.5))) | |
| ]) | |
| data_file = json.load( | |
| open(os.path.join(args.data_dir, "{}.{}.json".format(self.cur_la, 'train' if mode == 'train' else 'val')), | |
| 'r')) | |
| self.feature = self.load_data(data_file) | |
| def __len__(self): | |
| return len(self.feature['input_ids']) | |
| def __getitem__(self, index): | |
| input_ids = self.feature["input_ids"][index] | |
| # attention_mask = self.feature["attention_mask"][index] | |
| attention_mask = [1] * len(input_ids) | |
| labels = self.feature["labels"][index] | |
| bbox = self.feature["bbox"][index] | |
| segment_ids = self.feature['segment_ids'][index] | |
| position_ids = self.feature['position_ids'][index] | |
| img = pil_loader(self.feature['image_path'][index]) | |
| for_patches, _ = self.common_transform(img, augmentation=False) | |
| patch = self.patch_transform(for_patches) | |
| assert len(input_ids) == len(attention_mask) == len(labels) == len(bbox) == len(segment_ids) | |
| res = { | |
| "input_ids": input_ids, | |
| "attention_mask": attention_mask, | |
| "labels": labels, | |
| "bbox": bbox, | |
| "segment_ids": segment_ids, | |
| "position_ids": position_ids, | |
| "images": patch, | |
| } | |
| return res | |
| def pil_loader(path: str) -> Image.Image: | |
| # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835) | |
| with open(path, 'rb') as f: | |
| img = Image.open(f) | |
| return img.convert('RGB') | |