text_html_output_1.png
from datasets import Dataset

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()

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)