Spaces:
Runtime error
Runtime error
| 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 | |