Update files/functions.py

files/functions.py

@@ -70,36 +70,36 @@ label2color = {
 
 # bounding boxes start and end of a sequence
 cls_box = [0, 0, 0, 0]
+cls_box1, cls_box2 = cls_box, cls_box
+
 sep_box_lilt = cls_box
+sep_box1 = sep_box_lilt
+
 sep_box_layoutxlm = [1000, 1000, 1000, 1000]
+sep_box2 = sep_box_layoutxlm
 
 # models
 model_id_lilt = "pierreguillou/lilt-xlm-roberta-base-finetuned-with-DocLayNet-base-at-paragraphlevel-ml512"
+model_id1 = model_id_lilt
 model_id_layoutxlm = "pierreguillou/layout-xlm-base-finetuned-with-DocLayNet-base-at-paragraphlevel-ml512"
+model_id2 = model_id_layoutxlm
 
 # tokenizer for LayoutXLM
 tokenizer_id_layoutxlm = "xlm-roberta-base"
 
 # (tokenization) The maximum length of a feature (sequence)
-if str(384) in model_id_lilt:
-  max_length_lilt = 384 
-elif str(512) in model_id_lilt:
-  max_length_lilt = 512 
-else:
-  print("Error with max_length_lilt of chunks!")
-
-if str(384) in model_id_layoutxlm:
-  max_length_layoutxlm = 384 
-elif str(512) in model_id_layoutxlm:
-  max_length_layoutxlm = 512 
+if (str(384) in model_id_lilt) and (str(384) in model_id_layoutxlm):
+  max_length = 384 
+elif (str(512) in model_id_lilt) and (str(512) in model_id_layoutxlm):
+  max_length = 512 
 else:
-  print("Error with max_length_layoutxlm of chunks!")
+  print("Error with max_length of chunks!")
 
 # (tokenization) overlap
 doc_stride = 128 # The authorized overlap between two part of the context when splitting it is needed.
 
 # max PDF page images that will be displayed
-max_imgboxes = 1
+max_imgboxes = 2
 
 # get files
 examples_dir = 'files/'
@@ -159,6 +159,9 @@ tokenizer_lilt = AutoTokenizer.from_pretrained(model_id_lilt)
 model_lilt = AutoModelForTokenClassification.from_pretrained(model_id_lilt);
 model_lilt.to(device);
 
+tokenizer1 = tokenizer_lilt
+model1 = model_lilt
+
 ## model LayoutXLM
 from transformers import LayoutLMv2ForTokenClassification # LayoutXLMTokenizerFast, 
 model_layoutxlm = LayoutLMv2ForTokenClassification.from_pretrained(model_id_layoutxlm);
@@ -172,14 +175,8 @@ feature_extractor = LayoutLMv2FeatureExtractor(apply_ocr=False)
 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)
+tokenizer2 = tokenizer_layoutxlm
+model2 = model_layoutxlm
 
 
 # General
@@ -519,14 +516,10 @@ def extraction_data_from_image(images):
             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_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, texts_lines, texts_pars, texts_lines_par, row_indexes, par_boxes, line_boxes, lines_par_boxes
 
-            
-# Inference
-
 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()
@@ -711,8 +704,8 @@ def predictions_token_level(images, custom_encoded_dataset, model_id, model):
 
 from functools import reduce
 
-# Get predictions (line level)
-def predictions_paragraph_level(max_length, tokenizer, id2label, dataset, outputs, images_ids_list, chunk_ids, input_ids, bboxes, cls_box, sep_box):
+# Get predictions (paragraph level)
+def predictions_probs_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()
@@ -788,24 +781,69 @@ def predictions_paragraph_level(max_length, tokenizer, id2label, dataset, output
                     prob_label = reduce(lambda x, y: x*y, probs_list)
                     prob_label = prob_label**(1./(len(probs_list))) # normalization
                     probs_label.append(prob_label)
-                max_value = max(probs_label)
-                max_index = probs_label.index(max_value)
-                probs_bbox[str(bbox)] = max_index
+                # max_value = max(probs_label)
+                # max_index = probs_label.index(max_value)
+                # probs_bbox[str(bbox)] = max_index
+                probs_bbox[str(bbox)] = probs_label
 
             bboxes_list_dict[image_id] = bboxes_list
             input_ids_dict_dict[image_id] = input_ids_dict
             probs_dict_dict[image_id] = probs_bbox
 
+            # df[image_id] = pd.DataFrame()
+            # df[image_id]["bboxes"] = bboxes_list
+            # df[image_id]["texts"] = [tokenizer.decode(input_ids_dict[str(bbox)]) for bbox in bboxes_list]
+            # df[image_id]["labels"] = [id2label[probs_bbox[str(bbox)]] for bbox in bboxes_list]
+
+        return probs_bbox, bboxes_list_dict, input_ids_dict_dict, probs_dict_dict #, df
+
+    else:
+        print("An error occurred while getting predictions!")
+
+from functools import reduce
+
+# Get predictions (paragraph level)
+def predictions_paragraph_level(max_length, tokenizer1, id2label, dataset, outputs1, images_ids_list1, chunk_ids1, input_ids1, bboxes1, cls_box1, sep_box1, tokenizer2, outputs2, images_ids_list2, chunk_ids2, input_ids2, bboxes2, cls_box2, sep_box2):
+
+    bboxes_list_dict, input_ids_dict_dict, probs_dict_dict, df = dict(), dict(), dict(), dict()
+
+    probs_bbox1, bboxes_list_dict1, input_ids_dict_dict1, probs_dict_dict1 = predictions_probs_paragraph_level(max_length, tokenizer1, id2label, dataset, outputs1, images_ids_list1, chunk_ids1, input_ids1, bboxes1, cls_box1, sep_box1)
+    probs_bbox2, bboxes_list_dict2, input_ids_dict_dict2, probs_dict_dict2 = predictions_probs_paragraph_level(max_length, tokenizer2, id2label, dataset, outputs2, images_ids_list2, chunk_ids2, input_ids2, bboxes2, cls_box2, sep_box2)
+
+    if len(images_ids_list1) > 0:
+
+        for i, image_id in enumerate(images_ids_list1):
+
+            bboxes_list1 = bboxes_list_dict1[image_id]
+            input_ids_dict1 = input_ids_dict_dict1[image_id]
+            probs_bbox1 = probs_dict_dict1[image_id]
+
+            bboxes_list2 = bboxes_list_dict2[image_id]
+            input_ids_dict2 = input_ids_dict_dict2[image_id]
+            probs_bbox2 = probs_dict_dict2[image_id]
+
+            probs_bbox = dict()
+            for bbox in bboxes_list1:
+                prob_bbox = [(p1+p2)/2 for p1,p2 in zip(probs_bbox1[str(bbox)], probs_bbox2[str(bbox)])]
+                max_value = max(prob_bbox)
+                max_index = prob_bbox.index(max_value)
+                probs_bbox[str(bbox)] = max_index
+
+            bboxes_list_dict[image_id] = bboxes_list1
+            input_ids_dict_dict[image_id] = input_ids_dict1
+            probs_dict_dict[image_id] = probs_bbox
+
             df[image_id] = pd.DataFrame()
-            df[image_id]["bboxes"] = bboxes_list
-            df[image_id]["texts"] = [tokenizer.decode(input_ids_dict[str(bbox)]) for bbox in bboxes_list]
-            df[image_id]["labels"] = [id2label[probs_bbox[str(bbox)]] for bbox in bboxes_list]
+            df[image_id]["bboxes"] = bboxes_list1
+            df[image_id]["texts"] = [tokenizer1.decode(input_ids_dict1[str(bbox)]) for bbox in bboxes_list1]
+            df[image_id]["labels"] = [id2label[probs_bbox[str(bbox)]] for bbox in bboxes_list1]
 
-        return probs_bbox, bboxes_list_dict, input_ids_dict_dict, probs_dict_dict, df
+        return bboxes_list_dict, input_ids_dict_dict, probs_dict_dict, df
 
     else:
         print("An error occurred while getting predictions!")
 
+
 # Get labeled images with lines bounding boxes
 def get_labeled_images(id2label, dataset, images_ids_list, bboxes_list_dict, probs_dict_dict):
 
@@ -925,4 +963,5 @@ def display_chunk_lines_inference(dataset, encoded_dataset, index_chunk=None):
   print("\n>> Dataframe of annotated lines\n")
   cols = ["texts",	"bboxes"]
   df = df[cols]
-  display(df)
+  display(df)
+