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-DocLayNet-base_paragraphs_ml512"
+model_id = "pierreguillou/lilt-xlm-roberta-base-finetuned-with-DocLayNet-base-at-linelevel-ml384"
 
 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_paragraph(results, factor, conf_min=0):
+def get_data(results, factor, conf_min=0):
 
   data = {}
   for i in range(len(results['line_num'])):
@@ -160,55 +160,43 @@ def get_data_paragraph(results, factor, conf_min=0):
       par_idx += 1
 
   # get lines of texts, grouped by paragraph
-  texts_pars = list()
+  lines = 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])
-      texts_lines.append(line_text)
-      lines_par.append(line_text)
+      lines.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, lines_par_boxes = list(), list()
+  line_boxes = 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_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_boxes.append([int(xmin/factor), int(ymin/factor), int(xmax/factor), int(ymax/factor)])
     par_idx += 1
 
-  return texts_lines, texts_pars, texts_lines_par, row_indexes, par_boxes, line_boxes, lines_par_boxes 
+  return lines, row_indexes, par_boxes, line_boxes #data, par_data #
 
 # rescale image to get 300dpi
 def set_image_dpi_resize(image):
@@ -337,7 +325,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):
@@ -358,7 +346,7 @@ def pdf_to_images(uploaded_pdf):
         except PdfReadError:
             path_to_file = pdf_blank
             filename = path_to_file.replace(examples_dir,"")
-            msg = "invalid PDF file."
+            msg = "Invalid PDF file."
             images = [Image.open(image_blank)]
         else:
             try:
@@ -380,8 +368,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, 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()
+        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()
         
         try: 
             for i,image in enumerate(images):
@@ -393,7 +381,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()
@@ -413,40 +401,36 @@ def extraction_data_from_image(images):
                 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()
 
-                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])
+                lines[i], row_indexes[i], par_boxes[i], line_boxes[i] = get_data(results[i], factor, conf_min=0)
+                lines_list.append(lines[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_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})
+            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})
                     
             # print(f"The text data was successfully extracted by the OCR!")
 
-            return dataset, texts_lines, texts_pars, texts_lines_par, row_indexes, par_boxes, line_boxes, lines_par_boxes
+            return dataset, lines, row_indexes, par_boxes, line_boxes
 
 ## Inference
 
-def prepare_inference_features_paragraph(example, cls_box = cls_box, sep_box = sep_box):
+def prepare_inference_features(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()
 
-  # get batch
-  # batch_page_hash = example["page_hash"] 
+  # get batch 
   batch_images_ids = example["images_ids"]
   batch_images = example["images"]
-  batch_bboxes_par = example["bboxes_par"]
-  batch_texts_par = example["texts_par"]  
+  batch_bboxes_line = example["bboxes_line"]
+  batch_texts = example["texts"]  
   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]
@@ -455,37 +439,38 @@ def prepare_inference_features_paragraph(example, cls_box = cls_box, sep_box = s
   if not isinstance(batch_images_ids, list): 
     batch_images_ids = [batch_images_ids]
     batch_images = [batch_images]
-    batch_bboxes_par = [batch_bboxes_par]
-    batch_texts_par = [batch_texts_par]
+    batch_bboxes_line = [batch_bboxes_line]
+    batch_texts = [batch_texts]
     batch_width, batch_height = [batch_width], [batch_height] 
 
   # process all images of the batch
-  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)):
+  for num_batch, (image_id, boxes, texts, width, height) in enumerate(zip(batch_images_ids, batch_bboxes_line, batch_texts, batch_width, batch_height)):
     tokens_list = []
     bboxes_list = []
 
     # add a dimension if only on image
-    if not isinstance(texts_par, list):
-      texts_par, boxes = [texts_par], [boxes]
+    if not isinstance(texts, list):
+      texts, boxes = [texts], [boxes]
 
     # convert boxes to original
-    normalize_bboxes_par = [normalize_box(upperleft_to_lowerright(box), width, height) for box in boxes]
+    normalize_bboxes_line = [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_par = sort_data_wo_labels(normalize_bboxes_par, texts_par)
+    boxes, texts = sort_data_wo_labels(normalize_bboxes_line, texts)
 
     count = 0
-    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
+    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
 
     # 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_par), 
+    encodings = tokenizer(" ".join(texts), 
                           truncation=True,
                           padding="max_length", 
                           max_length=max_length, 
@@ -530,7 +515,7 @@ def prepare_inference_features_paragraph(example, cls_box = cls_box, sep_box = s
       "normalized_bboxes": bb_list,
   }
 
-  from torch.utils.data import Dataset
+from torch.utils.data import Dataset
 
 class CustomDataset(Dataset):
   def __init__(self, dataset, tokenizer):
@@ -552,7 +537,7 @@ class CustomDataset(Dataset):
     
     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):
@@ -561,7 +546,6 @@ 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):
@@ -605,7 +589,7 @@ def predictions_token_level(images, custom_encoded_dataset):
 from functools import reduce
 
 # Get predictions (line level)
-def predictions_paragraph_level_gradio(dataset, outputs, images_ids_list, chunk_ids, input_ids, bboxes):
+def predictions_line_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()
@@ -671,7 +655,7 @@ def predictions_paragraph_level_gradio(dataset, outputs, images_ids_list, chunk_
                         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)]
@@ -700,7 +684,7 @@ def predictions_paragraph_level_gradio(dataset, outputs, images_ids_list, chunk_
         print("An error occurred while getting predictions!")
 
 # Get labeled images with lines bounding boxes
-def get_labeled_images_gradio(dataset, images_ids_list, bboxes_list_dict, probs_dict_dict):
+def get_labeled_images(dataset, images_ids_list, bboxes_list_dict, probs_dict_dict):
 
     labeled_images = list()
 
@@ -784,7 +768,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_paragraphs_inference(index_chunk=None):
+def display_chunk_lines_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)
@@ -797,14 +781,14 @@ def display_chunk_paragraphs_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 paragraphs\n")
+  print(">> PDF image with bounding boxes of lines\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)))
@@ -815,7 +799,7 @@ def display_chunk_paragraphs_inference(index_chunk=None):
   cv2.waitKey(0)
 
   # display image dataframe
-  print("\n>> Dataframe of annotated paragraphs\n")
-  cols = ["texts",  "bboxes"]
+  print("\n>> Dataframe of annotated lines\n")
+  cols = ["texts",	"bboxes"]
   df = df[cols]
-  display(df)
+  display(df)