Update files/functions.py

files/functions.py

@@ -57,7 +57,7 @@ sep_box = cls_box
 from transformers import AutoTokenizer, AutoModelForTokenClassification
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-model_id = "pierreguillou/lilt-xlm-roberta-base-finetuned-with-DocLayNet-base-at-linelevel-ml384"
+model_id = "pierreguillou/lilt-xlm-roberta-base-finetuned-DocLayNet-base_paragraphs_ml512"
 
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForTokenClassification.from_pretrained(model_id);
@@ -117,7 +117,7 @@ langdetect2Tesseract = {v:k for k,v in Tesseract2langdetect.items()}
 # 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'])):
@@ -160,43 +160,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):
@@ -259,7 +271,8 @@ def original_box(box, original_width, original_height, coco_width, coco_height):
     ]
 
 def get_blocks(bboxes_block, categories, texts):
-    # get list of unique block boxes
+
+ # get list of unique block boxes
     bbox_block_dict, bboxes_block_list, bbox_block_prec = dict(), list(), list()
     for count_block, bbox_block in enumerate(bboxes_block):
       if bbox_block != bbox_block_prec:
@@ -324,7 +337,7 @@ 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):
@@ -367,8 +380,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 = dict(), dict(), dict(), dict(), dict()
-        images_ids_list, lines_list, par_boxes_list, line_boxes_list, images_list, page_no_list, num_pages_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 = 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, page_no_list, num_pages_list = list(), list(), list(), list(), list(), list(), list(), list(), list(), list()
         
         try: 
             for i,image in enumerate(images):
@@ -380,14 +393,13 @@ 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()
                 txt = re.sub(r" +", " ", txt) # multiple space
                 txt = re.sub(r"(\n\s*)+\n+", "\n", txt) # multiple line
                 # txt = os.popen(f'tesseract {img_filepath} - {custom_config}').read()
-
                 try:
                     langs = detect_langs(txt)
                     langs = [langdetect2Tesseract[langs[i].lang] for i in range(len(langs))]
@@ -395,36 +407,37 @@ def extraction_data_from_image(images):
                 except:
                     langs_string = "eng"
                 langs_string += '+osd'
-                print("langs_string", langs_string)
-                custom_config = f'--oem 3 --psm 3 -l {langs_string} tsv' # default config PyTesseract: --oem 3 --psm 3
+                custom_config = f'--oem 3 --psm 3 -l {langs_string}' # default config PyTesseract: --oem 3 --psm 3
 
                 # OCR PyTesseract | get data
                 results[i] = pytesseract.image_to_data(img, config=custom_config, output_type=pytesseract.Output.DICT)
                 # results[i] = os.popen(f'tesseract {img_filepath} - {custom_config}').read()
 
-                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_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, "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, "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!")
+            # 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(example):
+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 = list(), list(), list(), list(), list()
 
@@ -432,8 +445,8 @@ def prepare_inference_features(example):
   # batch_page_hash = example["page_hash"] 
   batch_images_ids = example["images_ids"]
   batch_images = example["images"]
-  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]
@@ -442,38 +455,37 @@ def prepare_inference_features(example):
   if not isinstance(batch_images_ids, list): 
     batch_images_ids = [batch_images_ids]
     batch_images = [batch_images]
-    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, boxes, texts, width, height) in enumerate(zip(batch_images_ids, batch_bboxes_line, batch_texts, batch_width, batch_height)):
+  for num_batch, (image_id, boxes, texts_par, width, height) in enumerate(zip(batch_images_ids, 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, 
@@ -518,7 +530,7 @@ def prepare_inference_features(example):
       "normalized_bboxes": bb_list,
   }
 
-from torch.utils.data import Dataset
+  from torch.utils.data import Dataset
 
 class CustomDataset(Dataset):
   def __init__(self, dataset, tokenizer):
@@ -537,11 +549,10 @@ class CustomDataset(Dataset):
     encoding["input_ids"] = example["input_ids"]
     encoding["attention_mask"] = example["attention_mask"]
     encoding["bbox"] = example["normalized_bboxes"]
-    # encoding["labels"] = example["labels"]
     
     return encoding
-      
-import torch.nn.functional as F
+
+    import torch.nn.functional as F
 
 # get predictions at token level
 def predictions_token_level(images, custom_encoded_dataset):
@@ -550,6 +561,7 @@ def predictions_token_level(images, custom_encoded_dataset):
     if num_imgs > 0:
 
         chunk_ids, input_ids, bboxes, outputs, token_predictions  = dict(), dict(), dict(), dict(), dict()
+        normalize_batch_bboxes_lines_pars = dict()
         images_ids_list = list()
 
         for i,encoding in enumerate(custom_encoded_dataset):
@@ -593,7 +605,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_gradio(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()
@@ -659,7 +671,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)]
@@ -688,7 +700,7 @@ def predictions_line_level(dataset, outputs, images_ids_list, chunk_ids, input_i
         print("An error occurred while getting predictions!")
 
 # Get labeled images with lines bounding boxes
-def get_labeled_images(dataset, images_ids_list, bboxes_list_dict, probs_dict_dict):
+def get_labeled_images_gradio(dataset, images_ids_list, bboxes_list_dict, probs_dict_dict):
 
     labeled_images = list()
 
@@ -772,7 +784,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)
@@ -785,7 +797,7 @@ 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()
@@ -803,7 +815,7 @@ def display_chunk_lines_inference(index_chunk=None):
   cv2.waitKey(0)
 
   # display image dataframe
-  print("\n>> Dataframe of annotated lines\n")
-  cols = ["texts",	"bboxes"]
+  print("\n>> Dataframe of annotated paragraphs\n")
+  cols = ["texts",  "bboxes"]
   df = df[cols]
   display(df)