Update files/functions.py

files/functions.py

@@ -98,7 +98,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-v2",
+        repo_id = "pierreguillou/Inference-APP-Document-Understanding-at-paragraphlevel-v2",
         filename = "files/" + file_name,
         repo_type = "space"
         )
@@ -140,10 +140,7 @@ langdetect2Tesseract = {v:k for k,v in Tesseract2langdetect.items()}
 import torch
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# model
-from transformers import LayoutLMv2ForTokenClassification
-
-model_id = "pierreguillou/layout-xlm-base-finetuned-with-DocLayNet-base-at-linelevel-ml384"
+from transformers import LayoutLMv2ForTokenClassification # LayoutXLMTokenizerFast, 
 
 model = LayoutLMv2ForTokenClassification.from_pretrained(model_id);
 model.to(device);
@@ -154,7 +151,6 @@ feature_extractor = LayoutLMv2FeatureExtractor(apply_ocr=False)
 
 # tokenizer
 from transformers import AutoTokenizer
-tokenizer_id = "xlm-roberta-base"
 tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
 
 # get labels
@@ -167,7 +163,7 @@ num_labels = len(id2label)
 # 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'])):
@@ -210,43 +206,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):
@@ -375,7 +383,7 @@ def sort_data_wo_labels(bboxes, texts):
     sorted_texts = np.array(texts, dtype=object)[sorted_bboxes_indexes].tolist()
 
     return sorted_bboxes, sorted_texts
-
+    
 ## PDF processing
 
 # get filename and images of PDF pages
@@ -419,8 +427,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):
@@ -432,7 +440,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()
@@ -455,38 +463,43 @@ 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
 
-def prepare_inference_features(example, cls_box = cls_box, sep_box = sep_box):
+def prepare_inference_features_paragraph(example, cls_box = cls_box, sep_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]
@@ -496,38 +509,37 @@ def prepare_inference_features(example, cls_box = cls_box, sep_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, 
@@ -654,7 +666,7 @@ def predictions_token_level(images, custom_encoded_dataset):
 from functools import reduce
 
 # Get predictions (line level)
-def predictions_line_level(dataset, outputs, images_ids_list, chunk_ids, input_ids, bboxes):
+def predictions_paragraph_level(dataset, outputs, images_ids_list, chunk_ids, input_ids, bboxes):
 
     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()
@@ -719,7 +731,7 @@ def predictions_line_level(dataset, outputs, images_ids_list, chunk_ids, input_i
                     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)]
@@ -832,7 +844,7 @@ def get_encoded_chunk_inference(index_chunk=None):
   return image, df, num_tokens, page_no, num_pages
 
 # display chunk of PDF image and its data
-def display_chunk_lines_inference(index_chunk=None):
+def display_chunk_paragraphs_inference(index_chunk=None):
 
   # get image and image data
   image, df, num_tokens, page_no, num_pages = get_encoded_chunk_inference(index_chunk=index_chunk)
@@ -845,14 +857,14 @@ def display_chunk_lines_inference(index_chunk=None):
   print(f'Chunk ({num_tokens} tokens) of the PDF (page: {page_no+1} / {num_pages})\n')
 
   # display image with bounding boxes
-  print(">> PDF image with bounding boxes of lines\n")
+  print(">> PDF image with bounding boxes of paragraphs\n")
   draw = ImageDraw.Draw(image)
             
   labels = list()
   for box, text in zip(bboxes, texts):
       color = "red"
       draw.rectangle(box, outline=color)
-  
+
   # resize image to original
   width, height = image.size
   image = image.resize((int(0.5*width), int(0.5*height)))
@@ -863,7 +875,7 @@ def display_chunk_lines_inference(index_chunk=None):
   cv2.waitKey(0)
 
   # display image dataframe
-  print("\n>> Dataframe of annotated lines\n")
+  print("\n>> Dataframe of annotated paragraphs\n")
   cols = ["texts",	"bboxes"]
   df = df[cols]
   display(df)