text_plain_output_2.png text_plain_output_1.png
from PIL import Image

from PIL import Image, ImageDraw, ImageFont

from PIL import Image, ImageDraw, ImageFont

from datasets import Dataset

from datasets import Features, Sequence, ClassLabel, Value, Array2D, Array3D

from transformers import LayoutLMv2Processor

import pandas as pd

df1 = pd.read_csv('../input/newspaper-articles-csv-ocr/ocr/df_valid.csv')
df1['imgdir'] = '../input/newspaper-articles/articles/validation'
df2 = pd.read_csv('../input/newspaper-articles-csv-ocr/ocr/df_train1.csv')
df2['imgdir'] = '../input/newspaper-articles/articles/training'
df3 = pd.read_csv('../input/newspaper-articles-csv-ocr/ocr/df_train2.csv')
df3['imgdir'] = '../input/newspaper-articles/articles/training'
df4 = pd.read_csv('../input/newspaper-articles-csv-ocr/ocr/df_train3.csv')
df4['imgdir'] = '../input/newspaper-articles/articles/training'
df = pd.concat([df1, df2, df3, df4])

df = df[df.text.apply(lambda x: len(str(x)) > 1)]
df = df.reset_index(drop=True)

labels = df['label'].unique()
labels = labels.tolist()

id2label = {v: k for v, k in enumerate(labels)}
label2id = {k: v for v, k in enumerate(labels)}

df.groupby(['label']).size()

label2color = {'title': 'Red', 'text': 'Green'}
def Draw_BBox(image_path, _bbox_labeled: pd.DataFrame):
    actual_boxes = []
    for idx, row in _bbox_labeled.iterrows():
        x0 = row['left']
        y0 = row['top']
        x1 = row['width'] + row['left']
        y1 = row['height'] + row['top']
        label = row['label']
        conf = row['conf']
        color = label2color[label]
        actual_box = [x0, y0, x1, y1]
        draw = ImageDraw.Draw(image, 'RGB')
        draw.rectangle(actual_box, outline=color, width=1)
        draw.text((actual_box[0] + 10, actual_box[1] - 10), text=f'{label}', fill=color)
    return image
sample_filename = df['image_id'].iloc[0]
sample_dir = df['imgdir'].iloc[0]
from PIL import Image, ImageDraw, ImageFont
image = Image.open(f'{sample_dir}/{sample_filename}')
image = image.convert('RGB').resize((1000, 1000))
sample_image_df = df[df['image_id'] == sample_filename]
sample_image_df = sample_image_df
Draw_BBox(image, sample_image_df)

df_meta = df
df_meta = df_meta.drop_duplicates(subset=['image_id'])
df_meta = df_meta.reset_index(drop=True)
df_meta['id'] = df_meta.index
df_meta['image_path'] = df_meta['imgdir'] + '/' + df_meta['image_id']
df_meta = df_meta[['image_id', 'image_path']]
df_meta['split'] = 'train'

ratio_split = int(len(df_meta) * 0.2)
df_meta.loc[df_meta.index[-ratio_split:], 'split'] = 'test'
df_meta

from datasets import Dataset
train_meta = df_meta[df_meta.split == 'train']
train_meta = train_meta.drop(columns=['split'])
train_meta = train_meta.reset_index(drop=True)
dataset_train = Dataset.from_pandas(train_meta)
test_meta = df_meta[df_meta.split == 'test']
test_meta = test_meta.drop(columns=['split'])
test_meta = test_meta.reset_index(drop=True)
dataset_test = Dataset.from_pandas(test_meta)

def get_words_and_boxes(examples):
    image_id = examples['image_id']
    record = {}
    objs1 = []
    objs2 = []
    objs3 = []
    for index2, row in df.query('image_id == @image_id').iterrows():
        class_id = label2id[row['label']]
        word = row['text']
        bbox_resized = [int(row['left']), int(row['top']), int(row['left']) + int(row['width']), int(row['top']) + int(row['height'])]
        obj1 = word
        obj2 = bbox_resized
        obj3 = class_id
        objs1.append(obj1)
        objs2.append(obj2)
        objs3.append(obj3)
        record['words'] = objs1
        record['bboxes'] = objs2
        record['classes'] = objs3
    examples['words'] = [objs1]
    examples['bboxes'] = [objs2]
    examples['classes'] = [objs3]
    return examples
dataset_train_with_bboxs = dataset_train.map(get_words_and_boxes, batched=True, batch_size=1)
dataset_test_with_bboxs = dataset_test.map(get_words_and_boxes, batched=True, batch_size=1)

from PIL import Image, ImageDraw, ImageFont
example = dataset_train_with_bboxs[0]
image = Image.open(example['image_path'])
image = image.convert('RGB')
image

from PIL import Image
from transformers import LayoutLMv2Processor
from datasets import Features, Sequence, ClassLabel, Value, Array2D, Array3D
from transformers import pipeline
processor = LayoutLMv2Processor.from_pretrained('microsoft/layoutlmv2-base-uncased', revision='no_ocr')
features = Features({'image': Array3D(dtype='int64', shape=(3, 224, 224)), 'input_ids': Sequence(feature=Value(dtype='int64')), 'attention_mask': Sequence(Value(dtype='int64')), 'token_type_ids': Sequence(Value(dtype='int64')), 'bbox': Array2D(dtype='int64', shape=(512, 4)), 'labels': Sequence(ClassLabel(names=labels))})
def preprocess_data(examples):
    images = [Image.open(path).convert('RGB') for path in examples['image_path']]
    words = examples['words']
    boxes = examples['bboxes']
    word_labels = examples['classes']
    encoded_inputs = processor(image, words, boxes=boxes, word_labels=word_labels, padding='max_length', truncation=True)
    return encoded_inputs
train_dataset = dataset_train_with_bboxs.map(preprocess_data, batched=True, batch_size=1, remove_columns=dataset_train_with_bboxs.column_names, features=features)
test_dataset = dataset_test_with_bboxs.map(preprocess_data, batched=True, batch_size=1, remove_columns=dataset_test_with_bboxs.column_names, features=features)

processor.tokenizer.decode(train_dataset['input_ids'][11])