Update files/functions.py

files/functions.py

@@ -51,22 +51,13 @@ label2color = {
 
 # bounding boxes start and end of a sequence
 cls_box = [0, 0, 0, 0]
-sep_box = cls_box
+sep_box = [1000, 1000, 1000, 1000]
 
 # model
-from transformers import AutoTokenizer, AutoModelForTokenClassification
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model_id = "pierreguillou/layout-xlm-base-finetuned-with-DocLayNet-base-at-linelevel-ml384"
 
-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);
-model.to(device);
-
-# get labels
-id2label = model.config.id2label
-label2id = model.config.label2id
-num_labels = len(id2label)
+# tokenizer
+tokenizer_id = "xlm-roberta-base"
 
 # (tokenization) The maximum length of a feature (sequence)
 if str(384) in model_id:
@@ -81,7 +72,21 @@ doc_stride = 128 # The authorized overlap between two part of the context when s
 
 # max PDF page images that will be displayed
 max_imgboxes = 2
+
+# get files
 examples_dir = 'files/'
+Path(examples_dir).mkdir(parents=True, exist_ok=True)
+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",
+        filename = "files/" + file_name,
+        repo_type = "space"
+        )
+    shutil.copy(path_to_file,examples_dir)
+
+# path to files
 image_wo_content = examples_dir + "wo_content.png" # image without content
 pdf_blank = examples_dir + "blank.pdf" # blank PDF
 image_blank = examples_dir + "blank.png" # blank image
@@ -368,8 +373,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, 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()
         
         try: 
             for i,image in enumerate(images):
@@ -401,11 +406,15 @@ 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()
 
+                # 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])
                 par_boxes_list.append(par_boxes[i])
                 line_boxes_list.append(line_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)
@@ -414,7 +423,7 @@ def extraction_data_from_image(images):
             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, "images_pixels": images_pixels_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!")
 
@@ -424,11 +433,12 @@ def extraction_data_from_image(images):
 
 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()
+  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 
   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_images_size = [image.size for image in batch_images]
@@ -439,12 +449,13 @@ def prepare_inference_features(example, cls_box = cls_box, sep_box = sep_box):
   if not isinstance(batch_images_ids, list): 
     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_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, image_pixels, boxes, texts, width, height) in enumerate(zip(batch_images_ids, batch_images_pixels, batch_bboxes_line, batch_texts, batch_width, batch_height)):
     tokens_list = []
     bboxes_list = []
 
@@ -506,6 +517,7 @@ def prepare_inference_features(example, cls_box = cls_box, sep_box = sep_box):
       bb_list.append(bb)
       images_ids_list.append(image_id)
       chunks_ids_list.append(i)
+      images_pixels_list.append(image_pixels)
         
   return {
       "images_ids": images_ids_list,
@@ -513,6 +525,7 @@ def prepare_inference_features(example, cls_box = cls_box, sep_box = sep_box):
       "input_ids": input_ids_list,
       "attention_mask": attention_mask_list,
       "normalized_bboxes": bb_list,
+      "images_pixels": images_pixels_list
   }
 
 from torch.utils.data import Dataset
@@ -534,18 +547,21 @@ class CustomDataset(Dataset):
     encoding["input_ids"] = example["input_ids"]
     encoding["attention_mask"] = example["attention_mask"]
     encoding["bbox"] = example["normalized_bboxes"]
+    encoding["images_pixels"] = example["images_pixels"]
     
     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):
 
     num_imgs = len(images)
     if num_imgs > 0:
 
-        chunk_ids, input_ids, bboxes, outputs, token_predictions  = dict(), dict(), dict(), dict(), dict()
+        chunk_ids, input_ids, bboxes, pixels_values, outputs, token_predictions  = dict(), dict(), dict(), dict(), dict(), dict()
         images_ids_list = list()
 
         for i,encoding in enumerate(custom_encoded_dataset):
@@ -556,6 +572,7 @@ def predictions_token_level(images, custom_encoded_dataset):
             input_id = torch.tensor(encoding['input_ids'])[None]
             attention_mask = torch.tensor(encoding['attention_mask'])[None]
             bbox = torch.tensor(encoding['bbox'])[None]
+            pixel_values = torch.tensor(encoding["images_pixels"])
 
             # save data in dictionnaries
             if image_id not in images_ids_list: images_ids_list.append(image_id)
@@ -569,14 +586,18 @@ def predictions_token_level(images, custom_encoded_dataset):
             if image_id in bboxes: bboxes[image_id].append(bbox)
             else: bboxes[image_id] = [bbox]
 
+            if image_id in pixels_values: pixels_values[image_id].append(pixel_values)
+            else: pixels_values[image_id] = [pixel_values]
+
             # get prediction with forward pass
             with torch.no_grad():
                 output = model(
-                    input_ids=input_id,
-                    attention_mask=attention_mask,
-                    bbox=bbox
+                    input_ids=input_id.to(device),
+                    attention_mask=attention_mask.to(device),
+                    bbox=bbox.to(device),
+                    image=pixel_values.to(device)
                     )
-              
+
             # save probabilities of predictions in dictionnary
             if image_id in outputs: outputs[image_id].append(F.softmax(output.logits.squeeze(), dim=-1))
             else: outputs[image_id] = [F.softmax(output.logits.squeeze(), dim=-1)]