Upload dataset.py

dataset.py

@@ -0,0 +1,127 @@
+# %%writefile dataset.py
+import pandas as pd
+import os
+from PIL import Image, ImageDraw
+import numpy as np
+import pytesseract
+import torch
+
+## I guess, I got my own script for it, from https://github.com/shabie/docformer/blob/master/src/docformer/dataset.py
+
+def rescale_bbox(bbox, img_width : int, 
+                 img_height : int, size : int = 1000):
+  x0, x1, y0, y1, width, height = bbox
+  x0 = int(size * (x0 / img_width))
+  x1 = int(size * (x1 / img_width))
+  y0 = int(size * (y0 / img_height))
+  y1 = int(size * (y1 / img_height))
+  width = int(size * (width / img_width))
+  height = int(size * (height / img_height))
+  return [x0, x1, y0, y1, width, height]
+
+def coordinate_features(df_row):
+    xmin, ymin, width, height = df_row["left"], df_row["top"], df_row["width"], df_row["height"]
+    return [xmin, xmin + width, ymin, ymin + height, width, height]  ## [xmin, xmax, ymin, ymax, width, height]
+
+def get_ocr_results(image_path : str):
+
+    """
+    Returns words and its bounding boxes from the image file path
+    image_path: string containing the path of the image
+    """
+
+    ## Getting the Image
+    image = Image.open(image_path)
+    width, height = image.size
+
+    ## OCR Processing
+    ocr_df = pytesseract.image_to_data(image, 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)
+    ocr_df = ocr_df.sort_values(by=['left', 'top'])  ## Sorting the values on the basis of left, top bounding box coordinates
+
+    ## Finally getting the words and the bounding box
+    words = list(ocr_df.text.apply(lambda x: str(x).strip()))
+    actual_bboxes = ocr_df.apply(coordinate_features, axis=1).values.tolist()
+
+    # add as extra columns
+    assert len(words) == len(actual_bboxes)
+    return {"words": words, "bbox": actual_bboxes}
+
+## Stealed from here: https://github.com/uakarsh/latr/blob/main/src/latr/dataset.py
+
+def get_tokens_with_boxes(unnormalized_word_boxes, list_of_words, 
+                          tokenizer, pad_token_id : int = 0, 
+                          pad_token_box = [0, 0, 0, 0, 0, 0], 
+                          max_seq_len = 512, 
+                          sep_token_box = [0, 0, 1000, 1000, 0, 0]
+                          ):
+    
+    '''
+    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 = []
+
+    ## CLS, SEP tokens have to be appended
+    unnormalized_token_boxes.extend([pad_token_box])
+    tokenized_words.extend([tokenizer.cls_token_id])  ## CLS Token Box is same as pad_token_box, if not, you can change here
+
+    ## Normal for loop
+    idx = 0
+    for box, word in zip(unnormalized_word_boxes, list_of_words):
+      if idx != 0:
+        new_word = " " + word
+      else:
+        new_word = word
+      current_tokens = tokenizer(new_word, add_special_tokens = False).input_ids
+      unnormalized_token_boxes.extend([box]*len(current_tokens))
+      tokenized_words.extend(current_tokens)
+      idx += 1
+
+    ## For post processing the token box
+    if len(unnormalized_token_boxes)<max_seq_len:
+      unnormalized_token_boxes.extend([sep_token_box])
+      unnormalized_token_boxes.extend([pad_token_box] * (max_seq_len-len(unnormalized_token_boxes)))
+    
+    else:
+      unnormalized_token_boxes[max_seq_len - 1] = sep_token_box
+
+    ## For post processing the tokenized words
+    if len(tokenized_words) < max_seq_len:
+      tokenized_words.extend([tokenizer.sep_token_id])
+      tokenized_words.extend([pad_token_id]* (max_seq_len-len(tokenized_words)))
+
+    else:
+      tokenized_words[max_seq_len - 1] = tokenizer.sep_token_id
+        
+    return unnormalized_token_boxes[:max_seq_len], tokenized_words[:max_seq_len]
+
+
+def create_features(
+    img_path : str,
+    tokenizer,
+    max_seq_length : int = 512,
+    size : int = 1000,
+    use_ocr : bool = True,
+    bounding_box = None,
+    words = None ):
+  
+  image = Image.open(img_path).convert("RGB")
+  ocr_results = get_ocr_results(img_path)
+  ocr_results['rescale_bbox'] = list(map(lambda x: rescale_bbox(x, *image.size, size  = size), ocr_results['bbox']))
+  boxes, words = get_tokens_with_boxes(ocr_results['rescale_bbox'], ocr_results['words'], tokenizer)
+  torch_boxes = torch.as_tensor(boxes)
+  torch_words = torch.as_tensor(words)
+
+  return torch_boxes, torch_words, ocr_results['bbox']