Update files/functions.py

files/functions.py

@@ -75,8 +75,8 @@ sep_box_lilt = cls_box
 sep_box_layoutxlm = [1000, 1000, 1000, 1000]
 
 # models
-model_id_lilt = "pierreguillou/lilt-xlm-roberta-base-finetuned-with-DocLayNet-base-at-linelevel-ml384"
-model_id_layoutxlm = "pierreguillou/layout-xlm-base-finetuned-with-DocLayNet-base-at-linelevel-ml384"
+model_id_lilt = "pierreguillou/lilt-xlm-roberta-base-finetuned-with-DocLayNet-base-at-paragraphlevel-ml512"
+model_id_layoutxlm = "pierreguillou/layout-xlm-base-finetuned-with-DocLayNet-base-at-paragraphlevel-ml512"
 
 # tokenizer for LayoutXLM
 tokenizer_id_layoutxlm = "xlm-roberta-base"
@@ -109,7 +109,7 @@ from huggingface_hub import hf_hub_download
 files = ["example.pdf", "blank.pdf", "blank.png", "languages_iso.csv", "languages_tesseract.csv", "wo_content.png"]
 for file_name in files:
     path_to_file = hf_hub_download(
-        repo_id = "pierreguillou/Inference-APP-Document-Understanding-at-linelevel-LiLT-base-LayoutXLM-base-v1",
+        repo_id = "pierreguillou/Inference-APP-Document-Understanding-at-paragraphlevel-LiLT-base-LayoutXLM-base-v1",
         filename = "files/" + file_name,
         repo_type = "space"
         )
@@ -146,47 +146,13 @@ for lang_t, langcode_t in zip(langs_t,langscode_t):
 
 langdetect2Tesseract = {v:k for k,v in Tesseract2langdetect.items()}
 
-## model / feature extractor / tokenizer
 
-# get device
-import torch
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-## model LiLT
-import transformers
-from transformers import AutoTokenizer, AutoModelForTokenClassification
-tokenizer_lilt = AutoTokenizer.from_pretrained(model_id_lilt)
-model_lilt = AutoModelForTokenClassification.from_pretrained(model_id_lilt);
-model_lilt.to(device);
-
-## model LayoutXLM
-from transformers import LayoutLMv2ForTokenClassification # LayoutXLMTokenizerFast, 
-model_layoutxlm = LayoutLMv2ForTokenClassification.from_pretrained(model_id_layoutxlm);
-model_layoutxlm.to(device);
-
-# feature extractor
-from transformers import LayoutLMv2FeatureExtractor
-feature_extractor = LayoutLMv2FeatureExtractor(apply_ocr=False)
-
-# tokenizer
-from transformers import AutoTokenizer
-tokenizer_layoutxlm = AutoTokenizer.from_pretrained(tokenizer_id_layoutxlm)
-
-# get labels
-id2label_lilt = model_lilt.config.id2label
-label2id_lilt = model_lilt.config.label2id
-num_labels_lilt = len(id2label_lilt)
-
-id2label_layoutxlm = model_layoutxlm.config.id2label
-label2id_layoutxlm = model_layoutxlm.config.label2id
-num_labels_layoutxlm = len(id2label_layoutxlm)
-
-## General
+# General
 
 # get text and bounding boxes from an image
 # https://stackoverflow.com/questions/61347755/how-can-i-get-line-coordinates-that-readed-by-tesseract
 # https://medium.com/geekculture/tesseract-ocr-understanding-the-contents-of-documents-beyond-their-text-a98704b7c655
-def get_data(results, factor, conf_min=0):
+def get_data_paragraph(results, factor, conf_min=0):
 
   data = {}
   for i in range(len(results['line_num'])):
@@ -229,43 +195,55 @@ def get_data(results, factor, conf_min=0):
       par_idx += 1
 
   # get lines of texts, grouped by paragraph
-  lines = list()
+  texts_pars = list()
   row_indexes = list()
+  texts_lines = list()
+  texts_lines_par = list()
   row_index = 0
   for _,par in par_data.items():
     count_lines = 0
+    lines_par = list()
     for _,line in par.items():
       if count_lines == 0: row_indexes.append(row_index)
       line_text = ' '.join([item[0] for item in line])
-      lines.append(line_text)
+      texts_lines.append(line_text)
+      lines_par.append(line_text)
       count_lines += 1
       row_index += 1
     # lines.append("\n")
     row_index += 1
+    texts_lines_par.append(lines_par)
+    texts_pars.append(' '.join(lines_par))
   # lines = lines[:-1]
   
   # get paragraphes boxes (par_boxes)
   # get lines boxes (line_boxes)
   par_boxes = list()
   par_idx = 1
-  line_boxes = list()
+  line_boxes, lines_par_boxes = list(), list()
   line_idx = 1
   for _, par in par_data.items():
     xmins, ymins, xmaxs, ymaxs = list(), list(), list(), list()
+    line_boxes_par = list()
+    count_line_par = 0
     for _, line in par.items():
       xmin, ymin = line[0][1], line[0][2]
       xmax, ymax = (line[-1][1] + line[-1][3]), (line[-1][2] + line[-1][4])
       line_boxes.append([int(xmin/factor), int(ymin/factor), int(xmax/factor), int(ymax/factor)])
+      line_boxes_par.append([int(xmin/factor), int(ymin/factor), int(xmax/factor), int(ymax/factor)])
       xmins.append(xmin)
       ymins.append(ymin)
       xmaxs.append(xmax)
       ymaxs.append(ymax)
       line_idx += 1
+      count_line_par += 1
     xmin, ymin, xmax, ymax = min(xmins), min(ymins), max(xmaxs), max(ymaxs)
-    par_boxes.append([int(xmin/factor), int(ymin/factor), int(xmax/factor), int(ymax/factor)])
+    par_bbox = [int(xmin/factor), int(ymin/factor), int(xmax/factor), int(ymax/factor)]
+    par_boxes.append(par_bbox)
+    lines_par_boxes.append(line_boxes_par)
     par_idx += 1
 
-  return lines, row_indexes, par_boxes, line_boxes #data, par_data #
+  return texts_lines, texts_pars, texts_lines_par, row_indexes, par_boxes, line_boxes, lines_par_boxes 
 
 # rescale image to get 300dpi
 def set_image_dpi_resize(image):
@@ -395,7 +373,8 @@ def sort_data_wo_labels(bboxes, texts):
 
     return sorted_bboxes, sorted_texts
 
-## PDF processing
+    
+# PDF processing
 
 # get filename and images of PDF pages
 def pdf_to_images(uploaded_pdf):
@@ -438,8 +417,8 @@ def extraction_data_from_image(images):
 
         # https://pyimagesearch.com/2021/11/15/tesseract-page-segmentation-modes-psms-explained-how-to-improve-your-ocr-accuracy/
         custom_config = r'--oem 3 --psm 3 -l eng' # default config PyTesseract: --oem 3 --psm 3 -l eng+deu+fra+jpn+por+spa+rus+hin+chi_sim
-        results, lines, row_indexes, par_boxes, line_boxes, images_pixels = dict(), dict(), dict(), dict(), dict(), dict()
-        images_ids_list, lines_list, par_boxes_list, line_boxes_list, images_list, images_pixels_list, page_no_list, num_pages_list = list(), list(), list(), list(), list(), list(), list(), list()
+        results, texts_lines, texts_pars, texts_lines_par, row_indexes, par_boxes, line_boxes, lines_par_boxes, images_pixels = dict(), dict(), dict(), dict(), dict(), dict(), dict(), dict(), dict()
+        images_ids_list, texts_lines_list, texts_pars_list, texts_lines_par_list, par_boxes_list, line_boxes_list, lines_par_boxes_list, images_list, images_pixels_list, page_no_list, num_pages_list = list(), list(), list(), list(), list(), list(), list(), list(), list(), list(), list()
         
         try: 
             for i,image in enumerate(images):
@@ -451,7 +430,7 @@ def extraction_data_from_image(images):
                 img = np.array(img, dtype='uint8') # convert PIL to cv2
                 img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # gray scale image
                 ret,img = cv2.threshold(img,127,255,cv2.THRESH_BINARY)
-
+                
                 # OCR PyTesseract | get langs of page
                 txt = pytesseract.image_to_string(img, config=custom_config)
                 txt = txt.strip().lower()
@@ -474,38 +453,44 @@ def extraction_data_from_image(images):
                 # get image pixels
                 images_pixels[i] = feature_extractor(images[i], return_tensors="pt").pixel_values
 
-                lines[i], row_indexes[i], par_boxes[i], line_boxes[i] = get_data(results[i], factor, conf_min=0)
-                lines_list.append(lines[i])
+                texts_lines[i], texts_pars[i], texts_lines_par[i], row_indexes[i], par_boxes[i], line_boxes[i], lines_par_boxes[i] = get_data_paragraph(results[i], factor, conf_min=0)
+                texts_lines_list.append(texts_lines[i])
+                texts_pars_list.append(texts_pars[i])
+                texts_lines_par_list.append(texts_lines_par[i])
                 par_boxes_list.append(par_boxes[i])
                 line_boxes_list.append(line_boxes[i])
+                lines_par_boxes_list.append(lines_par_boxes[i])
                 images_ids_list.append(i)
                 images_pixels_list.append(images_pixels[i])
                 images_list.append(images[i])
                 page_no_list.append(i)
-                num_pages_list.append(num_imgs)
+                num_pages_list.append(num_imgs) 
 
         except:
             print(f"There was an error within the extraction of PDF text by the OCR!")
         else: 
             from datasets import Dataset
-            dataset = Dataset.from_dict({"images_ids": images_ids_list, "images": images_list, "images_pixels": images_pixels_list, "page_no": page_no_list, "num_pages": num_pages_list, "texts": lines_list, "bboxes_line": line_boxes_list})
+            dataset = Dataset.from_dict({"images_ids": images_ids_list, "images": images_list, "images_pixels": images_pixels_list, "page_no": page_no_list, "num_pages": num_pages_list, "texts_line": texts_lines_list, "texts_par": texts_pars_list, "texts_lines_par": texts_lines_par_list, "bboxes_par": par_boxes_list, "bboxes_lines_par":lines_par_boxes_list})
                     
+
             # print(f"The text data was successfully extracted by the OCR!")
 
-            return dataset, lines, row_indexes, par_boxes, line_boxes
+            return dataset, texts_lines, texts_pars, texts_lines_par, row_indexes, par_boxes, line_boxes, lines_par_boxes
 
-## Inference
+            
+# Inference
 
-def prepare_inference_features(example, tokenizer, max_length, cls_box, sep_box):
+def prepare_inference_features_paragraph(example, tokenizer, max_length, cls_box, sep_box):
 
   images_ids_list, chunks_ids_list, input_ids_list, attention_mask_list, bb_list, images_pixels_list = list(), list(), list(), list(), list(), list()
 
-  # get batch 
+  # get batch
+  # batch_page_hash = example["page_hash"] 
   batch_images_ids = example["images_ids"]
   batch_images = example["images"]
   batch_images_pixels = example["images_pixels"]
-  batch_bboxes_line = example["bboxes_line"]
-  batch_texts = example["texts"]  
+  batch_bboxes_par = example["bboxes_par"]
+  batch_texts_par = example["texts_par"]  
   batch_images_size = [image.size for image in batch_images]
 
   batch_width, batch_height = [image_size[0] for image_size in batch_images_size], [image_size[1] for image_size in batch_images_size]
@@ -515,38 +500,37 @@ def prepare_inference_features(example, tokenizer, max_length, cls_box, sep_box)
     batch_images_ids = [batch_images_ids]
     batch_images = [batch_images]
     batch_images_pixels = [batch_images_pixels]
-    batch_bboxes_line = [batch_bboxes_line]
-    batch_texts = [batch_texts]
+    batch_bboxes_par = [batch_bboxes_par]
+    batch_texts_par = [batch_texts_par]
     batch_width, batch_height = [batch_width], [batch_height] 
 
   # process all images of the batch
-  for num_batch, (image_id, image_pixels, boxes, texts, width, height) in enumerate(zip(batch_images_ids, batch_images_pixels, batch_bboxes_line, batch_texts, batch_width, batch_height)):
+  for num_batch, (image_id, image_pixels, boxes, texts_par, width, height) in enumerate(zip(batch_images_ids, batch_images_pixels, batch_bboxes_par, batch_texts_par, batch_width, batch_height)):
     tokens_list = []
     bboxes_list = []
 
     # add a dimension if only on image
-    if not isinstance(texts, list):
-      texts, boxes = [texts], [boxes]
+    if not isinstance(texts_par, list):
+      texts_par, boxes = [texts_par], [boxes]
 
     # convert boxes to original
-    normalize_bboxes_line = [normalize_box(upperleft_to_lowerright(box), width, height) for box in boxes]
+    normalize_bboxes_par = [normalize_box(upperleft_to_lowerright(box), width, height) for box in boxes]
 
     # sort boxes with texts
     # we want sorted lists from top to bottom of the image
-    boxes, texts = sort_data_wo_labels(normalize_bboxes_line, texts)
+    boxes, texts_par = sort_data_wo_labels(normalize_bboxes_par, texts_par)
 
     count = 0
-    for box, text in zip(boxes, texts):
-      tokens = tokenizer.tokenize(text)
-      num_tokens = len(tokens) # get number of tokens
-      tokens_list.extend(tokens)
-
-      bboxes_list.extend([box] * num_tokens) # number of boxes must be the same as the number of tokens
+    for box, text_par in zip(boxes, texts_par):
+      tokens_par = tokenizer.tokenize(text_par)
+      num_tokens_par = len(tokens_par) # get number of tokens
+      tokens_list.extend(tokens_par)
+      bboxes_list.extend([box] * num_tokens_par) # number of boxes must be the same as the number of tokens
 
     # use of return_overflowing_tokens=True / stride=doc_stride
     # to get parts of image with overlap
     # source: https://huggingface.co/course/chapter6/3b?fw=tf#handling-long-contexts
-    encodings = tokenizer(" ".join(texts), 
+    encodings = tokenizer(" ".join(texts_par), 
                           truncation=True,
                           padding="max_length", 
                           max_length=max_length, 
@@ -681,7 +665,7 @@ def predictions_token_level(images, custom_encoded_dataset, model_id, model):
 from functools import reduce
 
 # Get predictions (line level)
-def predictions_line_level(max_length, tokenizer, id2label, dataset, outputs, images_ids_list, chunk_ids, input_ids, bboxes, cls_box, sep_box):
+def predictions_paragraph_level(max_length, tokenizer, id2label, dataset, outputs, images_ids_list, chunk_ids, input_ids, bboxes, cls_box, sep_box):
 
     ten_probs_dict, ten_input_ids_dict, ten_bboxes_dict = dict(), dict(), dict()
     bboxes_list_dict, input_ids_dict_dict, probs_dict_dict, df = dict(), dict(), dict(), dict()
@@ -746,7 +730,7 @@ def predictions_line_level(max_length, tokenizer, id2label, dataset, outputs, im
                     input_ids_dict[str(bbox)].append(input_id)
                     probs_dict[str(bbox)].append(probs)
                 bbox_prev = bbox
-
+                
             probs_bbox = dict()
             for i,bbox in enumerate(bboxes_list):
                 probs = probs_dict[str(bbox)]