Update files/functions.py

files/functions.py

@@ -178,3 +178,519 @@ id2label_layoutxlm = model_layoutxlm.config.id2label
 label2id_layoutxlm = model_layoutxlm.config.label2id
 num_labels_layoutxlm = len(id2label_layoutxlm)
 
+## PDf processing
+
+# get filename and images of PDF pages
+def pdf_to_images(uploaded_pdf):
+
+    # Check if None object
+    if uploaded_pdf is None:
+        path_to_file = pdf_blank
+        filename = path_to_file.replace(examples_dir,"")
+        msg = "Invalid PDF file."
+        images = [Image.open(image_blank)]
+    else:
+        # path to the uploaded PDF
+        path_to_file = uploaded_pdf.name
+        filename = path_to_file# .replace("/tmp/","")
+
+        try:
+            PdfReader(path_to_file)
+        except PdfReadError:
+            path_to_file = pdf_blank
+            filename = path_to_file.replace(examples_dir,"")
+            msg = "Invalid PDF file."
+            images = [Image.open(image_blank)]
+        else:
+            try:
+                # images = convert_from_path(path_to_file, last_page=max_imgboxes)
+
+                pdf = pdfium.PdfDocument(str(filename))
+                version = pdf.get_version()  # get the PDF standard version
+                n_pages = len(pdf)  # get the number of pages in the document
+                last_page = max_imgboxes
+                page_indices = [i for i in range(last_page)]  # pages until last_page
+                images = list(pdf.render(
+                    pdfium.PdfBitmap.to_pil,
+                    page_indices = page_indices,
+                    scale = 300/72,  # 300dpi resolution
+                ))
+                
+                num_imgs = len(images)
+                msg = f'The PDF "{filename}" was converted into {num_imgs} images.'
+            except:
+                msg = f'Error with the PDF "{filename}": it was not converted into images.'
+                images = [Image.open(image_wo_content)]
+
+    return filename, msg, images
+
+# Extraction of image data (text and bounding boxes)
+def extraction_data_from_image(images):
+
+    num_imgs = len(images)
+
+    if num_imgs > 0:
+
+        # 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, 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):
+                # image preprocessing
+                # https://docs.opencv.org/3.0-beta/doc/py_tutorials/py_imgproc/py_thresholding/py_thresholding.html
+                img = image.copy()
+                factor, path_to_img = set_image_dpi_resize(img) # Rescaling to 300dpi
+                img = Image.open(path_to_img)
+                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))]
+                    langs_string = '+'.join(langs)
+                except:
+                    langs_string = "eng"
+                langs_string += '+osd'
+                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()
+
+                # get image pixels
+                images_pixels[i] = feature_extractor(images[i], return_tensors="pt").pixel_values
+
+                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) 
+
+        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_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()
+
+  # 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_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]
+
+  # add a dimension if not a batch but only one image
+  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_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_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_par, list):
+      texts_par, boxes = [texts_par], [boxes]
+
+    # convert boxes to original
+    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_par = sort_data_wo_labels(normalize_bboxes_par, texts_par)
+
+    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
+
+    # 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), 
+                          truncation=True,
+                          padding="max_length", 
+                          max_length=max_length, 
+                          stride=doc_stride, 
+                          return_overflowing_tokens=True, 
+                          return_offsets_mapping=True
+                          )
+
+    otsm = encodings.pop("overflow_to_sample_mapping")
+    offset_mapping = encodings.pop("offset_mapping")
+
+    # Let's label those examples and get their boxes   
+    sequence_length_prev = 0   
+    for i, offsets in enumerate(offset_mapping):
+      # truncate tokens, boxes and labels based on length of chunk - 2 (special tokens <s> and </s>)
+      sequence_length = len(encodings.input_ids[i]) - 2
+      if i == 0: start = 0
+      else: start += sequence_length_prev - doc_stride
+      end = start + sequence_length
+      sequence_length_prev = sequence_length
+
+      # get tokens, boxes and labels of this image chunk
+      bb = [cls_box] + bboxes_list[start:end] + [sep_box]
+
+      # as the last chunk can have a length < max_length
+      # we must to add [tokenizer.pad_token] (tokens), [sep_box] (boxes) and [-100] (labels)
+      if len(bb) < max_length:
+        bb = bb + [sep_box] * (max_length - len(bb))
+
+      # append results
+      input_ids_list.append(encodings["input_ids"][i])
+      attention_mask_list.append(encodings["attention_mask"][i])
+      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,
+      "chunk_ids": chunks_ids_list,
+      "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
+
+class CustomDataset(Dataset):
+  def __init__(self, dataset, tokenizer):
+    self.dataset = dataset
+    self.tokenizer = tokenizer
+
+  def __len__(self):
+    return len(self.dataset)
+
+  def __getitem__(self, idx):
+    # get item
+    example = self.dataset[idx]
+    encoding = dict()
+    encoding["images_ids"] = example["images_ids"]
+    encoding["chunk_ids"] = example["chunk_ids"]
+    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
+
+# get predictions at token level
+def predictions_token_level(images, custom_encoded_dataset, model_id, model):
+
+    num_imgs = len(images)
+    if num_imgs > 0:
+
+        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):
+
+            # get custom encoded data
+            image_id = encoding['images_ids']
+            chunk_id = encoding['chunk_ids']
+            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)
+
+            if image_id in chunk_ids: chunk_ids[image_id].append(chunk_id)
+            else: chunk_ids[image_id] = [chunk_id]
+
+            if image_id in input_ids: input_ids[image_id].append(input_id)
+            else: input_ids[image_id] = [input_id]
+
+            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():
+
+                if model_id == model_id_lilt:
+                    output = model(
+                        input_ids=input_id.to(device),
+                        attention_mask=attention_mask.to(device),
+                        bbox=bbox.to(device),
+                        )
+                elif model_id == model_id_layoutxlm:
+                    output = model(
+                        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)]
+
+        return outputs, images_ids_list, chunk_ids, input_ids, bboxes
+
+    else:
+        print("An error occurred while getting predictions!")
+
+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):
+
+    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()
+
+    if len(images_ids_list) > 0:
+
+        for i, image_id in enumerate(images_ids_list):
+
+            # get image information
+            images_list = dataset.filter(lambda example: example["images_ids"] == image_id)["images"]
+            image = images_list[0]
+            width, height = image.size
+
+            # get data
+            chunk_ids_list = chunk_ids[image_id]
+            outputs_list = outputs[image_id]
+            input_ids_list = input_ids[image_id]
+            bboxes_list = bboxes[image_id]
+
+            # create zeros tensors
+            ten_probs = torch.zeros((outputs_list[0].shape[0] - 2)*len(outputs_list), outputs_list[0].shape[1])
+            ten_input_ids = torch.ones(size=(1, (outputs_list[0].shape[0] - 2)*len(outputs_list)), dtype =int)
+            ten_bboxes = torch.zeros(size=(1, (outputs_list[0].shape[0] - 2)*len(outputs_list), 4), dtype =int)
+
+            if len(outputs_list) > 1:
+              
+                for num_output, (output, input_id, bbox) in enumerate(zip(outputs_list, input_ids_list, bboxes_list)):
+                    start = num_output*(max_length - 2) - max(0,num_output)*doc_stride
+                    end = start + (max_length - 2)
+                    
+                    if num_output == 0:
+                        ten_probs[start:end,:] += output[1:-1]
+                        ten_input_ids[:,start:end] = input_id[:,1:-1]
+                        ten_bboxes[:,start:end,:] = bbox[:,1:-1,:]
+                    else:
+                        ten_probs[start:start + doc_stride,:] += output[1:1 + doc_stride]
+                        ten_probs[start:start + doc_stride,:] = ten_probs[start:start + doc_stride,:] * 0.5
+                        ten_probs[start + doc_stride:end,:] += output[1 + doc_stride:-1]
+
+                        ten_input_ids[:,start:start + doc_stride] = input_id[:,1:1 + doc_stride]
+                        ten_input_ids[:,start + doc_stride:end] = input_id[:,1 + doc_stride:-1]
+
+                        ten_bboxes[:,start:start + doc_stride,:] = bbox[:,1:1 + doc_stride,:]
+                        ten_bboxes[:,start + doc_stride:end,:] = bbox[:,1 + doc_stride:-1,:]
+              
+            else:
+                ten_probs += outputs_list[0][1:-1] 
+                ten_input_ids = input_ids_list[0][:,1:-1] 
+                ten_bboxes = bboxes_list[0][:,1:-1] 
+
+            ten_probs_list, ten_input_ids_list, ten_bboxes_list = ten_probs.tolist(), ten_input_ids.tolist()[0], ten_bboxes.tolist()[0]
+            bboxes_list = list()
+            input_ids_dict, probs_dict = dict(), dict()
+            bbox_prev = [-100, -100, -100, -100]
+            for probs, input_id, bbox in zip(ten_probs_list, ten_input_ids_list, ten_bboxes_list):
+                bbox = denormalize_box(bbox, width, height)
+                if bbox != bbox_prev and bbox != cls_box and bbox != sep_box and bbox[0] != bbox[2] and bbox[1] != bbox[3]:
+                    bboxes_list.append(bbox)
+                    input_ids_dict[str(bbox)] = [input_id]
+                    probs_dict[str(bbox)] = [probs]
+                elif bbox != cls_box and bbox != sep_box and bbox[0] != bbox[2] and bbox[1] != bbox[3]:
+                    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)]
+                probs = np.array(probs).T.tolist()
+            
+                probs_label = list()
+                for probs_list in probs:
+                    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
+
+            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!")
+
+# Get labeled images with lines bounding boxes
+def get_labeled_images(id2label, dataset, images_ids_list, bboxes_list_dict, probs_dict_dict):
+
+    labeled_images = list()
+
+    for i, image_id in enumerate(images_ids_list):
+
+        # get image
+        images_list = dataset.filter(lambda example: example["images_ids"] == image_id)["images"]
+        image = images_list[0]
+        width, height = image.size
+
+        # get predicted boxes and labels
+        bboxes_list = bboxes_list_dict[image_id]
+        probs_bbox = probs_dict_dict[image_id]
+
+        draw = ImageDraw.Draw(image)
+        # https://stackoverflow.com/questions/66274858/choosing-a-pil-imagefont-by-font-name-rather-than-filename-and-cross-platform-f
+        font = font_manager.FontProperties(family='sans-serif', weight='bold')
+        font_file = font_manager.findfont(font)
+        font_size = 30
+        font = ImageFont.truetype(font_file, font_size)
+
+        for bbox in bboxes_list:
+              predicted_label = id2label[probs_bbox[str(bbox)]]
+              draw.rectangle(bbox, outline=label2color[predicted_label])
+              draw.text((bbox[0] + 10, bbox[1] - font_size), text=predicted_label, fill=label2color[predicted_label], font=font)
+
+        labeled_images.append(image)
+
+    return labeled_images
+
+# get data of encoded chunk
+def get_encoded_chunk_inference(tokenizer, dataset, encoded_dataset, index_chunk=None):
+
+  # get datasets
+  example = dataset
+  encoded_example = encoded_dataset
+
+  # get randomly a document in dataset
+  if index_chunk == None: index_chunk = random.randint(0, len(encoded_example)-1)
+  encoded_example = encoded_example[index_chunk]
+  encoded_image_ids = encoded_example["images_ids"]
+
+  # get the image
+  example = example.filter(lambda example: example["images_ids"] ==  encoded_image_ids)[0]
+  image = example["images"] # original image
+  width, height = image.size
+  page_no = example["page_no"]
+  num_pages = example["num_pages"]
+
+  # get boxes, texts, categories
+  bboxes, input_ids  = encoded_example["normalized_bboxes"][1:-1], encoded_example["input_ids"][1:-1]
+  bboxes = [denormalize_box(bbox, width, height) for bbox in bboxes]
+  num_tokens = len(input_ids) + 2
+
+  # get unique bboxes and corresponding labels
+  bboxes_list, input_ids_list = list(), list()
+  input_ids_dict = dict()
+  bbox_prev = [-100, -100, -100, -100]
+  for i, (bbox, input_id) in enumerate(zip(bboxes, input_ids)):
+    if bbox != bbox_prev:
+      bboxes_list.append(bbox)
+      input_ids_dict[str(bbox)] = [input_id]
+    else:
+      input_ids_dict[str(bbox)].append(input_id)
+    
+    # start_indexes_list.append(i)
+    bbox_prev = bbox
+  
+  # do not keep "</s><pad><pad>..."
+  if input_ids_dict[str(bboxes_list[-1])][0] == (tokenizer.convert_tokens_to_ids('</s>')):
+    del input_ids_dict[str(bboxes_list[-1])]
+    bboxes_list = bboxes_list[:-1]
+
+  # get texts by line
+  input_ids_list = input_ids_dict.values()
+  texts_list = [tokenizer.decode(input_ids) for input_ids in input_ids_list]
+  
+  # display DataFrame
+  df = pd.DataFrame({"texts": texts_list, "input_ids": input_ids_list, "bboxes": bboxes_list})
+
+  return image, df, num_tokens, page_no, num_pages
+
+# display chunk of PDF image and its data
+def display_chunk_lines_inference(dataset, encoded_dataset, index_chunk=None):
+
+  # get image and image data
+  image, df, num_tokens, page_no, num_pages = get_encoded_chunk_inference(dataset, encoded_dataset, index_chunk=index_chunk)
+
+  # get data from dataframe
+  input_ids = df["input_ids"]
+  texts = df["texts"]
+  bboxes = df["bboxes"]
+
+  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")
+  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)))
+
+  # convert to cv and display
+  img = np.array(image, dtype='uint8') # PIL to cv2
+  cv2_imshow(img)
+  cv2.waitKey(0)
+
+  # display image dataframe
+  print("\n>> Dataframe of annotated lines\n")
+  cols = ["texts",	"bboxes"]
+  df = df[cols]
+  display(df)