EuroPython2022/latr-vqa

1import os2import json3import numpy as np4import pytesseract5from PIL import Image, ImageDraw6 7PAD_TOKEN_BOX = [0, 0, 0, 0]8max_seq_len = 5129 10## Function: 111## Purpose: Resize and align the bounding box for the different sized image12 13def resize_align_bbox(bbox, orig_w, orig_h, target_w, target_h):14    x_scale = target_w / orig_w15    y_scale = target_h / orig_h16    orig_left, orig_top, orig_right, orig_bottom = bbox17    x = int(np.round(orig_left * x_scale))18    y = int(np.round(orig_top * y_scale))19    xmax = int(np.round(orig_right * x_scale))20    ymax = int(np.round(orig_bottom * y_scale))21    return [x, y, xmax, ymax]22 23## Function: 224## Purpose: Reading the json file from the path and return the dictionary25 26def load_json_file(file_path):27  with open(file_path, 'r') as f:28    data = json.load(f)29  return data30 31## Function: 332## Purpose: Getting the address of specific file type, eg: .pdf, .tif, so and so33 34def get_specific_file(path, last_entry = 'tif'):35  base_path = path36  for i in os.listdir(path):37    if i.endswith(last_entry):38      return os.path.join(base_path, i)39 40  return '-1'41 42 43## Function: 444 45 46def 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):47    48    '''49    This function returns two items:50    1. unnormalized_token_boxes -> a list of len = max_seq_len, containing the boxes corresponding to the tokenized words, 51                                    one box might repeat as per the tokenization procedure52    2. tokenized_words -> tokenized words corresponding to the tokenizer and the list_of_words53    '''54 55    assert len(unnormalized_word_boxes) == len(list_of_words), "Bounding box length!= total words length"56    57    length_of_box = len(unnormalized_word_boxes)58    unnormalized_token_boxes = []59    tokenized_words = []60 61    for box, word in zip(unnormalized_word_boxes, list_of_words):62      current_tokens = tokenizer(word, add_special_tokens = False).input_ids63      unnormalized_token_boxes.extend([box]*len(current_tokens))64      tokenized_words.extend(current_tokens)65 66    if len(unnormalized_token_boxes)<max_seq_len:67        unnormalized_token_boxes.extend([pad_token_box] * (max_seq_len-len(unnormalized_token_boxes)))68        69    if len(tokenized_words)< max_seq_len:70        tokenized_words.extend([pad_token_id]* (max_seq_len-len(tokenized_words)))71        72    return unnormalized_token_boxes[:max_seq_len], tokenized_words[:max_seq_len]73 74## Function: 575## Function, which would only be used when the below function is used76 77def get_topleft_bottomright_coordinates(df_row):78    left, top, width, height = df_row["left"], df_row["top"], df_row["width"], df_row["height"]79    return [left, top, left + width, top + height]80 81## Function: 682## If the OCR is not provided, this function would help in extracting OCR83 84 85def apply_ocr(tif_path):86    """87    Returns words and its bounding boxes from an image88    """89    img = Image.open(tif_path).convert("RGB")90 91    ocr_df = pytesseract.image_to_data(img, output_type="data.frame")92    ocr_df = ocr_df.dropna().reset_index(drop=True)93    float_cols = ocr_df.select_dtypes("float").columns94    ocr_df[float_cols] = ocr_df[float_cols].round(0).astype(int)95    ocr_df = ocr_df.replace(r"^\s*$", np.nan, regex=True)96    ocr_df = ocr_df.dropna().reset_index(drop=True)97    words = list(ocr_df.text.apply(lambda x: str(x).strip()))98    actual_bboxes = ocr_df.apply(get_topleft_bottomright_coordinates, axis=1).values.tolist()99 100    # add as extra columns101    assert len(words) == len(actual_bboxes)102    return {"words": words, "bbox": actual_bboxes}103 104 105## Function: 7106## Merging all the above functions, for the purpose of extracting the image, bounding box and the tokens (sentence wise)107 108 109def create_features(110    image_path,111    tokenizer,112    target_size = (1000, 1000),113    max_seq_length=512,114    use_ocr = False,115    bounding_box = None,116    words = None117    ):118  119  '''120  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 ratio121  '''122 123 124  img = Image.open(image_path).convert("RGB")125  width_old, height_old = img.size126  img = img.resize(target_size)127  width, height = img.size128  129  ## Rescaling the bounding box as per the image size130  131 132  if (use_ocr == False) and (bounding_box == None or words == None):133    raise Exception('Please provide the bounding box and words or pass the argument "use_ocr" = True')134 135  if use_ocr == True:136    entries = apply_ocr(image_path)137    bounding_box = entries["bbox"]138    words = entries["words"]139  140  bounding_box = list(map(lambda x: resize_align_bbox(x,width_old,height_old, width, height), bounding_box))141  boxes, tokenized_words = get_tokens_with_boxes(unnormalized_word_boxes = bounding_box,142                                               list_of_words = words, 143                                               tokenizer = tokenizer,144                                               pad_token_id = 0,145                                               pad_token_box = PAD_TOKEN_BOX,146                                               max_seq_len = max_seq_length147                                               )148 149 150  return img, boxes, tokenized_words151