dataset.py

import os
import json
import numpy as np
import pytesseract
from PIL import Image, ImageDraw

PAD_TOKEN_BOX = [0, 0, 0, 0]
max_seq_len = 512

## Function: 1
## Purpose: Resize and align the bounding box for the different sized image

def resize_align_bbox(bbox, orig_w, orig_h, target_w, target_h):
    x_scale = target_w / orig_w
    y_scale = target_h / orig_h
    orig_left, orig_top, orig_right, orig_bottom = bbox
    x = int(np.round(orig_left * x_scale))
    y = int(np.round(orig_top * y_scale))
    xmax = int(np.round(orig_right * x_scale))
    ymax = int(np.round(orig_bottom * y_scale))
    return [x, y, xmax, ymax]

## Function: 2
## Purpose: Reading the json file from the path and return the dictionary

def load_json_file(file_path):
  with open(file_path, 'r') as f:
    data = json.load(f)
  return data

## Function: 3
## Purpose: Getting the address of specific file type, eg: .pdf, .tif, so and so

def get_specific_file(path, last_entry = 'tif'):
  base_path = path
  for i in os.listdir(path):
    if i.endswith(last_entry):
      return os.path.join(base_path, i)

  return '-1'


## Function: 4


def get_tokens_with_boxes(unnormalized_word_boxes, list_of_words, tokenizer, pad_token_id = 0, pad_token_box = [0, 0, 0, 0], max_seq_len = 512):
    
    '''
    This function returns two items:
    1. unnormalized_token_boxes -> a list of len = max_seq_len, containing the boxes corresponding to the tokenized words, 
                                    one box might repeat as per the tokenization procedure
    2. tokenized_words -> tokenized words corresponding to the tokenizer and the list_of_words
    '''

    assert len(unnormalized_word_boxes) == len(list_of_words), "Bounding box length!= total words length"
    
    length_of_box = len(unnormalized_word_boxes)
    unnormalized_token_boxes = []
    tokenized_words = []

    for box, word in zip(unnormalized_word_boxes, list_of_words):
      current_tokens = tokenizer(word, add_special_tokens = False).input_ids
      unnormalized_token_boxes.extend([box]*len(current_tokens))
      tokenized_words.extend(current_tokens)

    if len(unnormalized_token_boxes)<max_seq_len:
        unnormalized_token_boxes.extend([pad_token_box] * (max_seq_len-len(unnormalized_token_boxes)))
        
    if len(tokenized_words)< max_seq_len:
        tokenized_words.extend([pad_token_id]* (max_seq_len-len(tokenized_words)))
        
    return unnormalized_token_boxes[:max_seq_len], tokenized_words[:max_seq_len]

## Function: 5
## Function, which would only be used when the below function is used

def get_topleft_bottomright_coordinates(df_row):
    left, top, width, height = df_row["left"], df_row["top"], df_row["width"], df_row["height"]
    return [left, top, left + width, top + height]

## Function: 6
## If the OCR is not provided, this function would help in extracting OCR


def apply_ocr(tif_path):
    """
    Returns words and its bounding boxes from an image
    """
    img = Image.open(tif_path).convert("RGB")

    ocr_df = pytesseract.image_to_data(img, output_type="data.frame")
    ocr_df = ocr_df.dropna().reset_index(drop=True)
    float_cols = ocr_df.select_dtypes("float").columns
    ocr_df[float_cols] = ocr_df[float_cols].round(0).astype(int)
    ocr_df = ocr_df.replace(r"^\s*$", np.nan, regex=True)
    ocr_df = ocr_df.dropna().reset_index(drop=True)
    words = list(ocr_df.text.apply(lambda x: str(x).strip()))
    actual_bboxes = ocr_df.apply(get_topleft_bottomright_coordinates, axis=1).values.tolist()

    # add as extra columns
    assert len(words) == len(actual_bboxes)
    return {"words": words, "bbox": actual_bboxes}


## Function: 7
## Merging all the above functions, for the purpose of extracting the image, bounding box and the tokens (sentence wise)


def create_features(
    image_path,
    tokenizer,
    target_size = (1000, 1000),
    max_seq_length=512,
    use_ocr = False,
    bounding_box = None,
    words = None
    ):
  
  '''
  We assume that the bounding box provided are given as per the image scale (i.e not normalized), so that we just need to scale it as per the ratio
  '''


  img = Image.open(image_path).convert("RGB")
  width_old, height_old = img.size
  img = img.resize(target_size)
  width, height = img.size
  
  ## Rescaling the bounding box as per the image size
  

  if (use_ocr == False) and (bounding_box == None or words == None):
    raise Exception('Please provide the bounding box and words or pass the argument "use_ocr" = True')

  if use_ocr == True:
    entries = apply_ocr(image_path)
    bounding_box = entries["bbox"]
    words = entries["words"]
  
  bounding_box = list(map(lambda x: resize_align_bbox(x,width_old,height_old, width, height), bounding_box))
  boxes, tokenized_words = get_tokens_with_boxes(unnormalized_word_boxes = bounding_box,
                                               list_of_words = words, 
                                               tokenizer = tokenizer,
                                               pad_token_id = 0,
                                               pad_token_box = PAD_TOKEN_BOX,
                                               max_seq_len = max_seq_length
                                               )


  return img, boxes, tokenized_words