AuditEdge's picture
initial commit
81e13bb
raw
history blame
6.2 kB
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')