Update files/functions.py

files/functions.py

@@ -178,6 +178,233 @@ id2label_layoutxlm = model_layoutxlm.config.id2label
 label2id_layoutxlm = model_layoutxlm.config.label2id
 num_labels_layoutxlm = len(id2label_layoutxlm)
 
+## General
+
+# 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):
+
+  data = {}
+  for i in range(len(results['line_num'])):
+    level = results['level'][i]
+    block_num = results['block_num'][i]
+    par_num = results['par_num'][i]
+    line_num = results['line_num'][i]
+    top, left = results['top'][i], results['left'][i]
+    width, height = results['width'][i], results['height'][i]
+    conf = results['conf'][i]
+    text = results['text'][i]
+    if not (text == '' or text.isspace()):
+      if conf >= conf_min:
+        tup = (text, left, top, width, height)
+        if block_num in list(data.keys()):
+          if par_num in list(data[block_num].keys()):
+            if line_num in list(data[block_num][par_num].keys()):
+              data[block_num][par_num][line_num].append(tup)
+            else:
+              data[block_num][par_num][line_num] = [tup]
+          else:
+            data[block_num][par_num] = {}
+            data[block_num][par_num][line_num] = [tup]
+        else:
+            data[block_num] = {}
+            data[block_num][par_num] = {}
+            data[block_num][par_num][line_num] = [tup]
+
+  # get paragraphs dicionnary with list of lines
+  par_data = {}
+  par_idx = 1
+  for _, b  in data.items():
+    for _, p in b.items():
+      line_data = {}
+      line_idx = 1
+      for _, l in p.items():
+        line_data[line_idx] = l
+        line_idx += 1
+      par_data[par_idx] = line_data 
+      par_idx += 1
+
+  # get lines of texts, grouped by paragraph
+  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])
+      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, 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_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 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):
+    """
+    Rescaling image to 300dpi while resizing
+    :param image: An image
+    :return: A rescaled image
+    """
+    length_x, width_y = image.size
+    factor = min(1, float(1024.0 / length_x))
+    size = int(factor * length_x), int(factor * width_y)
+    # image_resize = image.resize(size, Image.Resampling.LANCZOS)
+    image_resize = image.resize(size, Image.LANCZOS)
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='1.png')
+    temp_filename = temp_file.name
+    image_resize.save(temp_filename, dpi=(300, 300))
+    return factor, temp_filename
+
+# it is important that each bounding box should be in (upper left, lower right) format.
+# source: https://github.com/NielsRogge/Transformers-Tutorials/issues/129
+def upperleft_to_lowerright(bbox):
+  x0, y0, x1, y1 = tuple(bbox)
+  if bbox[2] < bbox[0]:
+    x0 = bbox[2]
+    x1 = bbox[0] 
+  if bbox[3] < bbox[1]:
+    y0 = bbox[3]
+    y1 = bbox[1] 
+  return [x0, y0, x1, y1]
+
+# convert boundings boxes (left, top, width, height) format to (left, top, left+widght, top+height) format. 
+def convert_box(bbox):
+    x, y, w, h = tuple(bbox) # the row comes in (left, top, width, height) format
+    return [x, y, x+w, y+h] # we turn it into (left, top, left+widght, top+height) to get the actual box 
+
+# LiLT model gets 1000x10000 pixels images
+def normalize_box(bbox, width, height):
+    return [
+        int(1000 * (bbox[0] / width)),
+        int(1000 * (bbox[1] / height)),
+        int(1000 * (bbox[2] / width)),
+        int(1000 * (bbox[3] / height)),
+    ]
+
+# LiLT model gets 1000x10000 pixels images
+def denormalize_box(bbox, width, height):
+    return [
+        int(width * (bbox[0] / 1000)),
+        int(height * (bbox[1] / 1000)),
+        int(width* (bbox[2] / 1000)),
+        int(height * (bbox[3] / 1000)),
+    ]
+
+# get back original size
+def original_box(box, original_width, original_height, coco_width, coco_height):
+    return [
+        int(original_width * (box[0] / coco_width)),
+        int(original_height * (box[1] / coco_height)),
+        int(original_width * (box[2] / coco_width)),
+        int(original_height* (box[3] / coco_height)),
+    ]
+
+def get_blocks(bboxes_block, categories, texts):
+
+ # 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:
+        bbox_block_indexes = [i for i, bbox in enumerate(bboxes_block) if bbox == bbox_block]
+        bbox_block_dict[count_block] = bbox_block_indexes
+        bboxes_block_list.append(bbox_block)
+      bbox_block_prec = bbox_block
+
+    # get list of categories and texts by unique block boxes
+    category_block_list, text_block_list = list(), list()
+    for bbox_block in bboxes_block_list:
+      count_block = bboxes_block.index(bbox_block)
+      bbox_block_indexes = bbox_block_dict[count_block]
+      category_block = np.array(categories, dtype=object)[bbox_block_indexes].tolist()[0]
+      category_block_list.append(category_block)
+      text_block = np.array(texts, dtype=object)[bbox_block_indexes].tolist()
+      text_block = [text.replace("\n","").strip() for text in text_block]
+      if id2label[category_block] == "Text" or id2label[category_block] == "Caption" or id2label[category_block] == "Footnote":
+        text_block = ' '.join(text_block)
+      else:
+        text_block = '\n'.join(text_block)
+      text_block_list.append(text_block)
+
+    return bboxes_block_list, category_block_list, text_block_list
+
+# function to sort bounding boxes
+def get_sorted_boxes(bboxes):
+
+  # sort by y from page top to bottom 
+  sorted_bboxes = sorted(bboxes, key=itemgetter(1), reverse=False)
+  y_list = [bbox[1] for bbox in sorted_bboxes]
+
+  # sort by x from page left to right when boxes with same y
+  if len(list(set(y_list))) != len(y_list):
+    y_list_duplicates_indexes = dict()
+    y_list_duplicates = [item for item, count in collections.Counter(y_list).items() if count > 1]
+    for item in y_list_duplicates:
+      y_list_duplicates_indexes[item] = [i for i, e in enumerate(y_list) if e == item]
+      bbox_list_y_duplicates = sorted(np.array(sorted_bboxes, dtype=object)[y_list_duplicates_indexes[item]].tolist(), key=itemgetter(0), reverse=False)
+      np_array_bboxes = np.array(sorted_bboxes)
+      np_array_bboxes[y_list_duplicates_indexes[item]] = np.array(bbox_list_y_duplicates)
+      sorted_bboxes = np_array_bboxes.tolist()
+
+  return sorted_bboxes
+
+# sort data from y = 0 to end of page (and after, x=0 to end of page when necessary)
+def sort_data(bboxes, categories, texts):
+
+    sorted_bboxes = get_sorted_boxes(bboxes)
+    sorted_bboxes_indexes = [bboxes.index(bbox) for bbox in sorted_bboxes]
+    sorted_categories = np.array(categories, dtype=object)[sorted_bboxes_indexes].tolist()
+    sorted_texts = np.array(texts, dtype=object)[sorted_bboxes_indexes].tolist()
+
+    return sorted_bboxes, sorted_categories, sorted_texts
+
+# sort data from y = 0 to end of page (and after, x=0 to end of page when necessary)
+def sort_data_wo_labels(bboxes, texts):
+
+    sorted_bboxes = get_sorted_boxes(bboxes)
+    sorted_bboxes_indexes = [bboxes.index(bbox) for bbox in sorted_bboxes]
+    sorted_texts = np.array(texts, dtype=object)[sorted_bboxes_indexes].tolist()
+
+    return sorted_bboxes, sorted_texts
+
+    
 ## PDf processing
 
 # get filename and images of PDF pages