Delete TDTSR.py

TDTSR.py

@@ -1,332 +0,0 @@
-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