Spaces:
Running
Running
File size: 6,198 Bytes
81e13bb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 |
import pandas as pd
import numpy as np
import os
import argparse
from datasets.features import ClassLabel
from transformers import AutoProcessor
from sklearn.model_selection import train_test_split
from datasets import Features, Sequence, ClassLabel, Value, Array2D, Array3D, Dataset
from datasets import Image as Img
from PIL import Image
import warnings
warnings.filterwarnings('ignore')
def read_text_file(file_path):
with open(file_path, 'r') as f:
return (f.readlines())
def prepare_examples(examples):
images = examples[image_column_name]
words = examples[text_column_name]
boxes = examples[boxes_column_name]
word_labels = examples[label_column_name]
encoding = processor(images, words, boxes=boxes, word_labels=word_labels,
truncation=True, padding="max_length")
return encoding
def get_zip_dir_name():
try:
os.chdir('/content/data')
dir_list = os.listdir()
any_file_name = dir_list[0]
zip_dir_name = any_file_name[:any_file_name.find('\\')]
if all(list(map(lambda x: x.startswith(zip_dir_name), dir_list))):
return zip_dir_name
return False
finally:
os.chdir('./../')
def filter_out_unannotated(example):
tags = example['ner_tags']
return not all([tag == label2id['O'] for tag in tags])
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--valid_size')
parser.add_argument('--output_path')
args = parser.parse_args()
TEST_SIZE = float(args.valid_size)
OUTPUT_PATH = args.output_path
os.makedirs(args.output_path, exist_ok=True)
files = {}
zip_dir_name = get_zip_dir_name()
if zip_dir_name:
files['train_box'] = read_text_file(os.path.join(
os.curdir, 'data', f'{zip_dir_name}\\{zip_dir_name}_box.txt'))
files['train_image'] = read_text_file(os.path.join(
os.curdir, 'data', f'{zip_dir_name}\\{zip_dir_name}_image.txt'))
files['train'] = read_text_file(os.path.join(
os.curdir, 'data', f'{zip_dir_name}\\{zip_dir_name}.txt'))
else:
for f in os.listdir():
if f.endswith('.txt') and f.find('box') != -1:
files['train_box'] = read_text_file(os.path.join(os.curdir, f))
elif f.endswith('.txt') and f.find('image') != -1:
files['train_image'] = read_text_file(
os.path.join(os.curdir, f))
elif f.endswith('.txt') and f.find('labels') == -1:
files['train'] = read_text_file(os.path.join(os.curdir, f))
assert(len(files['train']) == len(files['train_box']))
assert(len(files['train_box']) == len(files['train_image']))
assert(len(files['train_image']) == len(files['train']))
images = {}
for i, row in enumerate(files['train_image']):
if row != '\n':
image_name = row.split('\t')[-1]
images.setdefault(image_name.replace('\n', ''), []).append(i)
words, bboxes, ner_tags, image_path = [], [], [], []
for image, rows in images.items():
words.append([row.split('\t')[0].replace('\n', '')
for row in files['train'][rows[0]:rows[-1]+1]])
ner_tags.append([row.split('\t')[1].replace('\n', '')
for row in files['train'][rows[0]:rows[-1]+1]])
bboxes.append([box.split('\t')[1].replace('\n', '')
for box in files['train_box'][rows[0]:rows[-1]+1]])
if zip_dir_name:
image_path.append(f"/content/data/{zip_dir_name}\\{image}")
else:
image_path.append(f"/content/data/{image}")
labels = list(set([tag for doc_tag in ner_tags for tag in doc_tag]))
id2label = {v: k for v, k in enumerate(labels)}
label2id = {k: v for v, k in enumerate(labels)}
dataset_dict = {
'id': range(len(words)),
'tokens': words,
'bboxes': [[list(map(int, bbox.split())) for bbox in doc] for doc in bboxes],
'ner_tags': [[label2id[tag] for tag in ner_tag] for ner_tag in ner_tags],
'image': [Image.open(path).convert("RGB") for path in image_path]
}
#raw features
features = Features({
'id': Value(dtype='string', id=None),
'tokens': Sequence(feature=Value(dtype='string', id=None), length=-1, id=None),
'bboxes': Sequence(feature=Sequence(feature=Value(dtype='int64', id=None), length=-1, id=None), length=-1, id=None),
'ner_tags': Sequence(feature=ClassLabel(num_classes=len(labels), names=labels, names_file=None, id=None), length=-1, id=None),
'image': Img(decode=True, id=None)
})
full_data_set = Dataset.from_dict(dataset_dict, features=features)
dataset = full_data_set.train_test_split(test_size=TEST_SIZE)
dataset["train"] = dataset["train"].filter(filter_out_unannotated)
processor = AutoProcessor.from_pretrained(
"microsoft/layoutlmv3-base", apply_ocr=False)
features = dataset["train"].features
column_names = dataset["train"].column_names
image_column_name = "image"
text_column_name = "tokens"
boxes_column_name = "bboxes"
label_column_name = "ner_tags"
# we need to define custom features for `set_format` (used later on) to work properly
features = Features({
'pixel_values': Array3D(dtype="float32", shape=(3, 224, 224)),
'input_ids': Sequence(feature=Value(dtype='int64')),
'attention_mask': Sequence(Value(dtype='int64')),
'bbox': Array2D(dtype="int64", shape=(512, 4)),
'labels': Sequence(ClassLabel(names=labels)),
})
train_dataset = dataset["train"].map(
prepare_examples,
batched=True,
remove_columns=column_names,
features=features,
)
eval_dataset = dataset["test"].map(
prepare_examples,
batched=True,
remove_columns=column_names,
features=features,
)
train_dataset.set_format("torch")
if not OUTPUT_PATH.endswith('/'):
OUTPUT_PATH += '/'
train_dataset.save_to_disk(f'{OUTPUT_PATH}train_split')
eval_dataset.save_to_disk(f'{OUTPUT_PATH}eval_split')
dataset.save_to_disk(f'{OUTPUT_PATH}raw_data')
|