Upload TDTSR.py

TDTSR.py

@@ -0,0 +1,332 @@
+import os
+import cv2
+from transformers import DetrFeatureExtractor
+from transformers import DetrForObjectDetection
+import torch
+import matplotlib.pyplot as plt
+from matplotlib.patches import Circle, Wedge, Rectangle
+import streamlit as st
+from PIL import Image
+import math
+
+
+colors = ["red", "blue", "green", "yellow", "orange", "violet"]
+
+
+def table_detector(image, THRESHOLD_PROBA):
+  '''
+  Table detection using DEtect-object TRansformer pre-trained on 1 million tables
+  
+  '''
+
+  feature_extractor = DetrFeatureExtractor(do_resize=True, size=800, max_size=800)
+  encoding = feature_extractor(image, return_tensors="pt")
+  # encoding.keys()
+  model = DetrForObjectDetection.from_pretrained("SalML/DETR-table-detection")
+  # SalML\DETR-table-detection
+  with torch.no_grad():
+    outputs = model(**encoding)
+
+  # keep only predictions of queries with 0.9+ confidence (excluding no-object class)
+  probas = outputs.logits.softmax(-1)[0, :, :-1]
+  keep = probas.max(-1).values > THRESHOLD_PROBA
+
+  # rescale bounding boxes
+  target_sizes = torch.tensor(image.size[::-1]).unsqueeze(0)
+  postprocessed_outputs = feature_extractor.post_process(outputs, target_sizes)
+  bboxes_scaled = postprocessed_outputs[0]['boxes'][keep]
+
+  return (model, image, probas[keep], bboxes_scaled)
+
+
+def table_struct_recog(image, THRESHOLD_PROBA):
+  '''
+  Table structure recognition using DEtect-object TRansformer pre-trained on 1 million tables
+  '''
+
+  feature_extractor = DetrFeatureExtractor(do_resize=True, size=1000, max_size=1000)
+  encoding = feature_extractor(image, return_tensors="pt")
+
+  model = DetrForObjectDetection.from_pretrained("SalML/DETR-table-structure-recognition")
+  with torch.no_grad():
+    outputs = model(**encoding)
+
+  # keep only predictions of queries with 0.9+ confidence (excluding no-object class)
+  probas = outputs.logits.softmax(-1)[0, :, :-1]
+  keep = probas.max(-1).values > THRESHOLD_PROBA
+
+  # rescale bounding boxes
+  target_sizes = torch.tensor(image.size[::-1]).unsqueeze(0)
+  postprocessed_outputs = feature_extractor.post_process(outputs, target_sizes)
+  bboxes_scaled = postprocessed_outputs[0]['boxes'][keep]
+
+  return (model, image, probas[keep], bboxes_scaled)
+
+def add_margin(pil_img, top=20, right=20, bottom=20, left=20, color=(255,255,255)):
+  '''
+  Image padding as part of TSR pre-processing to prevent missing table edges
+  '''
+  width, height = pil_img.size
+  new_width = width + right + left
+  new_height = height + top + bottom
+  result = Image.new(pil_img.mode, (new_width, new_height), color)
+  result.paste(pil_img, (left, top))
+  return result
+
+def plot_results_detection(c1, model, pil_img, prob, boxes, show_only_cropped=False):
+  '''
+  Plots the full pillow pdf-page image and adds a rectangle patch for table detection
+  '''
+
+  plt.figure(figsize=(32,20))
+  plt.imshow(pil_img)
+  ax = plt.gca()
+
+  for p, (xmin, ymin, xmax, ymax) in zip(prob, boxes.tolist()):
+
+      cl = p.argmax()
+      xmin, ymin, xmax, ymax = xmin-3, ymin-3, xmax+3, ymax+3  
+      ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,fill=False, color=colors[cl.item()], linewidth=3))
+      text = f'{model.config.id2label[cl.item()]}: {p[cl]:0.2f}'
+      ax.text(xmin, ymin, text, fontsize=15,bbox=dict(facecolor='yellow', alpha=0.5))
+  plt.axis('off')
+  plt.show()
+  c1.pyplot()
+
+
+def plot_table_detection(c2, model, pil_img, prob, boxes):
+  '''
+  Plots only the cropped table(s) from the table detection 
+  '''
+
+  plt.figure(figsize=(32,20))
+  ax = plt.gca()
+  cropped_img_list = []
+
+  for p, (xmin, ymin, xmax, ymax) in zip(prob, boxes.tolist()):
+
+      xmin, ymin, xmax, ymax = xmin-3, ymin-3, xmax+3, ymax+3  
+      cropped_img = pil_img.crop((xmin, ymin, xmax, ymax))
+      cropped_img_list.append(cropped_img)
+
+  for cropped_img in cropped_img_list:
+    plt.imshow(cropped_img)
+
+    plt.axis('off')
+    plt.show()
+    c2.pyplot()
+  return cropped_img_list
+
+
+def plot_structure(c3, model, pil_img, prob, boxes, class_to_show=0):
+  '''
+  To plot table pillow image and the TSR bounding boxes on the table
+  '''
+  plt.figure(figsize=(32,20))
+  plt.imshow(pil_img)
+  ax = plt.gca()
+  rows = {}
+  cols = {}
+  header = {}
+  row_header = {}
+  idx = 0
+
+  for p, (xmin, ymin, xmax, ymax) in zip(prob, boxes.tolist()):
+
+      xmin, ymin, xmax, ymax = xmin-3, ymin-3, xmax+3, ymax+3  
+      cl = p.argmax()
+      class_text = model.config.id2label[cl.item()]
+      text = f'{class_text}: {p[cl]:0.2f}'
+      # st.write(class_text)
+      if class_text != 'table':
+        
+        ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,fill=False, color=colors[cl.item()], linewidth=3))
+        ax.text(xmin, ymin, text, fontsize=15, bbox=dict(facecolor='yellow', alpha=0.5))
+
+        # if class_text == 'table column header':
+        #   header['header'] = (xmin, ymin, xmax, ymax)
+        if class_text == 'table row':
+          rows['table row '+str(idx)] = (xmin, ymin, xmax, ymax)
+        if class_text == 'table column':
+          cols['table column '+str(idx)] = (xmin, ymin, xmax, ymax)
+        # if class_text == 'table projected row header':
+        #   row_header['header table row'+str(idx)] = (xmin, ymin, xmax, ymax)
+
+      idx += 1
+
+  plt.show()
+  c3.pyplot()
+  # return header, row_header, rows, cols
+  return rows, cols
+
+
+
+def sort_table_features(header, row_header, rows, cols):
+  # Sometimes the header and first row overlap, and we need the header bbox not to have first row's bbox inside the headers bbox
+  y_header = header['header'][3] - 10
+  rows_ = {table_feature : (xmin, ymin, xmax, ymax) for table_feature, (xmin, ymin, xmax, ymax) in sorted(rows.items(), key=lambda tup: tup[1][1]) if ymin > y_header}
+  cols_ = {table_feature : (xmin, ymin, xmax, ymax) for table_feature, (xmin, ymin, xmax, ymax) in sorted(cols.items(), key=lambda tup: tup[1][0])}
+
+  row_header_ = {table_feature : (xmin, ymin, xmax, ymax) for table_feature, (xmin, ymin, xmax, ymax) in sorted(row_header.items(), key=lambda tup: tup[1][1])}
+
+  new_row = {}
+  idx = 0
+
+  for k1, v1 in rows_.items():
+    save_row = True
+    row_xmin, row_ymin, row_xmax, row_ymax = v1
+    for k2, v2 in row_header_.items():
+      header_row_xmin, header_row_ymin, header_row_xmax, header_row_ymax = v2
+      # table row and header table row are within 2 pixel range, skip saving the row
+      if math.isclose(row_ymin, header_row_ymin, abs_tol=2):
+        save_row = False
+    if save_row:
+      new_row['table row.'+str(idx)] = (row_xmin, row_ymin, row_xmax, row_ymax)
+      idx += 1
+
+  new_row_ = {table_feature : (xmin, ymin, xmax, ymax) for table_feature, (xmin, ymin, xmax, ymax) in sorted(new_row.items(), key=lambda tup: tup[1][1])}
+
+  return row_header_, new_row_, cols_
+
+
+def sort_table_featuresv2(rows, cols):
+  # Sometimes the header and first row overlap, and we need the header bbox not to have first row's bbox inside the headers bbox
+  rows_ = {table_feature : (xmin, ymin, xmax, ymax) for table_feature, (xmin, ymin, xmax, ymax) in sorted(rows.items(), key=lambda tup: tup[1][1])}
+  cols_ = {table_feature : (xmin, ymin, xmax, ymax) for table_feature, (xmin, ymin, xmax, ymax) in sorted(cols.items(), key=lambda tup: tup[1][0])}
+
+  return rows_, cols_
+
+def individual_table_features(pil_img, header, row_header, rows, cols):
+
+  for k, v in header.items():
+    xmin, ymin, xmax, ymax = v
+    cropped_img = pil_img.crop((xmin, ymin, xmax, ymax))
+    header[k] = xmin, ymin, xmax, ymax, cropped_img
+
+  for k, v in row_header.items():
+    xmin, ymin, xmax, ymax = v
+    cropped_img = pil_img.crop((xmin, ymin, xmax, ymax))
+    row_header[k] = xmin, ymin, xmax, ymax, cropped_img
+
+  for k, v in rows.items():
+    xmin, ymin, xmax, ymax = v
+    cropped_img = pil_img.crop((xmin, ymin, xmax, ymax))
+    rows[k] = xmin, ymin, xmax, ymax, cropped_img
+
+
+  for k, v in cols.items():
+    xmin, ymin, xmax, ymax = v
+    cropped_img = pil_img.crop((xmin, ymin, xmax, ymax))
+    cols[k] = xmin, ymin, xmax, ymax, cropped_img
+
+  return header, row_header, rows, cols
+
+def individual_table_featuresv2(pil_img, rows, cols):
+
+
+  for k, v in rows.items():
+    xmin, ymin, xmax, ymax = v
+    cropped_img = pil_img.crop((xmin, ymin, xmax, ymax))
+    rows[k] = xmin, ymin, xmax, ymax, cropped_img
+
+
+  for k, v in cols.items():
+    xmin, ymin, xmax, ymax = v
+    cropped_img = pil_img.crop((xmin, ymin, xmax, ymax))
+    cols[k] = xmin, ymin, xmax, ymax, cropped_img
+
+  return rows, cols
+
+def plot_table_features(c2, header, row_header, rows, cols):
+
+  for k, v in header.items():
+    _, _, _, _, pil_img = v
+
+  for k, v in row_header.items():
+    _, _, _, _, pil_img = v
+
+  for k, v in rows.items():
+    _, _, _, _, pil_img = v
+
+  for k, v in cols.items():
+    _, _, _, _, pil_img = v
+
+
+def master_row_set(header, row_header, rows, cols):
+  master_row = {**header, **row_header, **rows}
+  master_row_ = {table_feature : (xmin, ymin, xmax, ymax, img) for table_feature, (xmin, ymin, xmax, ymax, img) in sorted(master_row.items(), key=lambda tup: tup[1][1])}
+
+  return master_row_
+
+
+
+
+def object_to_cells(master_row, cols):
+  '''
+  Iterates to every row, be it header/simple row/header table row, cuts rows into cells and saves images in dictionary where length of dictionary = total rows
+  '''
+  cells_img = {}
+  header_idx = 0
+  row_idx = 0
+  for k_row, v_row in master_row.items():
+
+    if k_row[:16] == 'header table row':
+
+      _, _, _, _, row_header_img = v_row
+      cells_img[k_row+'.'+str(row_idx)] = row_header_img
+      row_idx += 1
+
+    elif k_row == 'header':
+
+      _, ymin, _, ymax, header_img = v_row
+
+      xa, ya, xb, yb = 0, 0, 0, ymax-ymin
+      for k_col, v_col in cols.items():
+        xmin_col, _, xmax_col, _, col_img = v_col
+        xa = xmin_col-19
+        xb = xmax_col-20
+
+        header_img_cropped = header_img.crop((xa, ya, xb, yb))
+        cells_img[k_row+'.'+str(header_idx)] = header_img_cropped
+        header_idx += 1
+
+
+    elif k_row[:9] == 'table row':
+
+      xmin, ymin, xmax, ymax, row_img = v_row
+      xa, ya, xb, yb = 0, 0, 0, ymax-ymin
+      row_img_list = []
+      for k_col, v_col in cols.items():
+        xmin_col, _, xmax_col, _, col_img = v_col
+        xa = xmin_col-19
+        xb = xmax_col-20
+        row_img_cropped = row_img.crop((xa, ya, xb, yb))
+        row_img_list.append(row_img_cropped)
+      cells_img[k_row+'.'+str(row_idx)] = row_img_list
+      row_idx += 1
+    
+  return cells_img
+
+
+def object_to_cellsv2(master_row, cols):
+  '''
+  Iterates to every row, be it header/simple row/header table row, cuts rows into cells and saves images in dictionary where length of dictionary = total rows
+  '''
+  cells_img = {}
+  header_idx = 0
+  row_idx = 0
+  for k_row, v_row in master_row.items():
+
+    xmin, ymin, xmax, ymax, row_img = v_row
+    xa, ya, xb, yb = 0, 0, 0, ymax-ymin
+    row_img_list = []
+    for k_col, v_col in cols.items():
+      xmin_col, _, xmax_col, _, col_img = v_col
+      xa = xmin_col-19
+      xb = xmax_col-20
+      row_img_cropped = row_img.crop((xa, ya, xb, yb))
+      row_img_list.append(row_img_cropped)
+    cells_img[k_row+'.'+str(row_idx)] = row_img_list
+    row_idx += 1
+    
+  return cells_img