Update app.py

app.py

@@ -1,11 +1,12 @@
 from PIL import Image, ImageEnhance, ImageOps
 import string
-from collections import Counter 
+from collections import Counter
 from itertools import tee, count
 import pytesseract
 from pytesseract import Output
 import json
 import pandas as pd
+
 # import matplotlib.pyplot as plt
 import cv2
 import numpy as np
@@ -14,35 +15,65 @@ from transformers import TableTransformerForObjectDetection
 import torch
 import gradio as gr
 
-def plot_results_detection(model, image, prob, bboxes_scaled,  delta_xmin, delta_ymin, delta_xmax, delta_ymax):
-  plt.imshow(image)
-  ax = plt.gca()
 
-  for p, (xmin, ymin, xmax, ymax) in zip(prob, bboxes_scaled.tolist()):
-      cl = p.argmax()
-      xmin, ymin, xmax, ymax = xmin-delta_xmin, ymin-delta_ymin, xmax+delta_xmax, ymax+delta_ymax 
-      ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,fill=False, color='red', linewidth=3))
-      text = f'{model.config.id2label[cl.item()]}: {p[cl]:0.2f}'
-      ax.text(xmin-20, ymin-50, text, fontsize=10,bbox=dict(facecolor='yellow', alpha=0.5))
-  plt.axis('off')
+def plot_results_detection(
+    model, image, prob, bboxes_scaled, delta_xmin, delta_ymin, delta_xmax, delta_ymax
+):
+    plt.imshow(image)
+    ax = plt.gca()
+
+    for p, (xmin, ymin, xmax, ymax) in zip(prob, bboxes_scaled.tolist()):
+        cl = p.argmax()
+        xmin, ymin, xmax, ymax = (
+            xmin - delta_xmin,
+            ymin - delta_ymin,
+            xmax + delta_xmax,
+            ymax + delta_ymax,
+        )
+        ax.add_patch(
+            plt.Rectangle(
+                (xmin, ymin),
+                xmax - xmin,
+                ymax - ymin,
+                fill=False,
+                color="red",
+                linewidth=3,
+            )
+        )
+        text = f"{model.config.id2label[cl.item()]}: {p[cl]:0.2f}"
+        ax.text(
+            xmin - 20,
+            ymin - 50,
+            text,
+            fontsize=10,
+            bbox=dict(facecolor="yellow", alpha=0.5),
+        )
+    plt.axis("off")
+
 
 def crop_tables(pil_img, prob, boxes, delta_xmin, delta_ymin, delta_xmax, delta_ymax):
-    '''
-    crop_tables and plot_results_detection must have same co-ord shifts because 1 only plots the other one updates co-ordinates 
-    '''
+    """
+    crop_tables and plot_results_detection must have same co-ord shifts because 1 only plots the other one updates co-ordinates
+    """
     cropped_img_list = []
 
     for p, (xmin, ymin, xmax, ymax) in zip(prob, boxes.tolist()):
 
-        xmin, ymin, xmax, ymax = xmin-delta_xmin, ymin-delta_ymin, xmax+delta_xmax, ymax+delta_ymax 
+        xmin, ymin, xmax, ymax = (
+            xmin - delta_xmin,
+            ymin - delta_ymin,
+            xmax + delta_xmax,
+            ymax + delta_ymax,
+        )
         cropped_img = pil_img.crop((xmin, ymin, xmax, ymax))
         cropped_img_list.append(cropped_img)
     return cropped_img_list
 
-def add_padding(pil_img, top, right, bottom, left, color=(255,255,255)):
-    '''
+
+def add_padding(pil_img, top, right, bottom, left, 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
@@ -50,15 +81,18 @@ def add_padding(pil_img, top, right, bottom, left, color=(255,255,255)):
     result.paste(pil_img, (left, top))
     return result
 
+
 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")
 
-    model = TableTransformerForObjectDetection.from_pretrained("microsoft/table-transformer-detection")
+    model = TableTransformerForObjectDetection.from_pretrained(
+        "microsoft/table-transformer-detection"
+    )
 
     with torch.no_grad():
         outputs = model(**encoding)
@@ -68,20 +102,22 @@ def table_detector(image, THRESHOLD_PROBA):
 
     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]
+    bboxes_scaled = postprocessed_outputs[0]["boxes"][keep]
 
     return (model, 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 = TableTransformerForObjectDetection.from_pretrained("microsoft/table-transformer-structure-recognition")
+    model = TableTransformerForObjectDetection.from_pretrained(
+        "microsoft/table-transformer-structure-recognition"
+    )
     with torch.no_grad():
         outputs = model(**encoding)
 
@@ -90,16 +126,19 @@ def table_struct_recog(image, THRESHOLD_PROBA):
 
     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]
+    bboxes_scaled = postprocessed_outputs[0]["boxes"][keep]
 
     return (model, probas[keep], bboxes_scaled)
 
-def generate_structure(model, pil_img, prob, boxes, expand_rowcol_bbox_top, expand_rowcol_bbox_bottom):
+
+def generate_structure(
+    model, pil_img, prob, boxes, expand_rowcol_bbox_top, expand_rowcol_bbox_bottom
+):
     colors = ["red", "blue", "green", "yellow", "orange", "violet"]
-    '''
+    """
     Co-ordinates are adjusted here by 3 'pixels'
     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()
@@ -108,33 +147,55 @@ def generate_structure(model, pil_img, prob, boxes, expand_rowcol_bbox_top, expa
     idx = 0
     for p, (xmin, ymin, xmax, ymax) in zip(prob, boxes.tolist()):
 
-        xmin, ymin, xmax, ymax = xmin, ymin, xmax, ymax 
+        xmin, ymin, xmax, ymax = xmin, ymin, xmax, ymax
         cl = p.argmax()
         class_text = model.config.id2label[cl.item()]
-        text = f'{class_text}: {p[cl]:0.2f}'
+        text = f"{class_text}: {p[cl]:0.2f}"
         # or (class_text == 'table column')
         # if (class_text == 'table row')  or (class_text =='table projected row header') or (class_text == 'table column'):
-            # ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,fill=False, color=colors[0], linewidth=2))
-            # ax.text(xmin-10, ymin-10, text, fontsize=5, bbox=dict(facecolor='yellow', alpha=0.5))
-
-        if class_text == 'table row':
-            rows['table row.'+str(idx)] = (xmin, ymin-expand_rowcol_bbox_top, xmax, ymax+expand_rowcol_bbox_bottom)
-        if class_text == 'table column':
-            cols['table column.'+str(idx)] = (xmin, ymin-expand_rowcol_bbox_top, xmax, ymax+expand_rowcol_bbox_bottom)
+        # ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,fill=False, color=colors[0], linewidth=2))
+        # ax.text(xmin-10, ymin-10, text, fontsize=5, bbox=dict(facecolor='yellow', alpha=0.5))
+
+        if class_text == "table row":
+            rows["table row." + str(idx)] = (
+                xmin,
+                ymin - expand_rowcol_bbox_top,
+                xmax,
+                ymax + expand_rowcol_bbox_bottom,
+            )
+        if class_text == "table column":
+            cols["table column." + str(idx)] = (
+                xmin,
+                ymin - expand_rowcol_bbox_top,
+                xmax,
+                ymax + expand_rowcol_bbox_bottom,
+            )
 
         idx += 1
 
     # plt.axis('on')
     return rows, cols
 
-def sort_table_featuresv2(rows:dict, cols:dict):
+
+def sort_table_featuresv2(rows: dict, cols: dict):
     # 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])}
+    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_featuresv2(pil_img, rows:dict, cols:dict):
+
+def individual_table_featuresv2(pil_img, rows: dict, cols: dict):
 
     for k, v in rows.items():
         xmin, ymin, xmax, ymax = v
@@ -148,12 +209,19 @@ def individual_table_featuresv2(pil_img, rows:dict, cols:dict):
 
     return rows, cols
 
-def object_to_cellsv2(master_row:dict, cols:dict, expand_rowcol_bbox_top, expand_rowcol_bbox_bottom, padd_left):
-    '''Removes redundant bbox for rows&columns and divides each row into cells from columns
+
+def object_to_cellsv2(
+    master_row: dict,
+    cols: dict,
+    expand_rowcol_bbox_top,
+    expand_rowcol_bbox_bottom,
+    padd_left,
+):
+    """Removes redundant bbox for rows&columns and divides each row into cells from columns
     Args:
     Returns:
-    
-    '''
+
+    """
     cells_img = {}
     header_idx = 0
     row_idx = 0
@@ -181,7 +249,7 @@ def object_to_cellsv2(master_row:dict, cols:dict, expand_rowcol_bbox_top, expand
     #     new_master_row[k_row] = v_row
     ######################################################
     for k_row, v_row in new_master_row.items():
-        
+
         _, _, _, _, row_img = v_row
         xmax, ymax = row_img.size
         xa, ya, xb, yb = 0, 0, 0, ymax
@@ -201,31 +269,39 @@ def object_to_cellsv2(master_row:dict, cols:dict, expand_rowcol_bbox_top, expand
             xb = xmax_col
             if idx == 0:
                 xa = 0
-            if idx == len(new_cols)-1:
+            if idx == len(new_cols) - 1:
                 xb = xmax
             xa, ya, xb, yb = xa, ya, xb, yb
 
             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
+        cells_img[k_row + "." + str(row_idx)] = row_img_list
         row_idx += 1
 
-    return cells_img, len(new_cols), len(new_master_row)-1
+    return cells_img, len(new_cols), len(new_master_row) - 1
+
 
 def pytess(cell_pil_img):
-    return ' '.join(pytesseract.image_to_data(cell_pil_img, output_type=Output.DICT, config='-c tessedit_char_blacklist=œ˜â€œï¬â™Ã©œ¢!|”?«“¥ --psm 6 preserve_interword_spaces')['text']).strip()
+    return " ".join(
+        pytesseract.image_to_data(
+            cell_pil_img,
+            output_type=Output.DICT,
+            config="-c tessedit_char_blacklist=œ˜â€œï¬â™Ã©œ¢!|”?«“¥ --psm 6 preserve_interword_spaces",
+        )["text"]
+    ).strip()
 
-def uniquify(seq, suffs = count(1)):
+
+def uniquify(seq, suffs=count(1)):
     """Make all the items unique by adding a suffix (1, 2, etc).
     Credit: https://stackoverflow.com/questions/30650474/python-rename-duplicates-in-list-with-progressive-numbers-without-sorting-list
     `seq` is mutable sequence of strings.
     `suffs` is an optional alternative suffix iterable.
     """
-    not_unique = [k for k,v in Counter(seq).items() if v>1] 
+    not_unique = [k for k, v in Counter(seq).items() if v > 1]
 
-    suff_gens = dict(zip(not_unique, tee(suffs, len(not_unique))))  
-    for idx,s in enumerate(seq):
+    suff_gens = dict(zip(not_unique, tee(suffs, len(not_unique))))
+    for idx, s in enumerate(seq):
         try:
             suffix = str(next(suff_gens[s]))
         except KeyError:
@@ -235,34 +311,36 @@ def uniquify(seq, suffs = count(1)):
 
     return seq
 
+
 def clean_dataframe(df):
-    '''
+    """
     Remove irrelevant symbols that appear with tesseractOCR
-    '''
+    """
     # df.columns = [col.replace('|', '') for col in df.columns]
 
     for col in df.columns:
 
-        df[col]=df[col].str.replace("'", '', regex=True)
-        df[col]=df[col].str.replace('"', '', regex=True)
-        df[col]=df[col].str.replace(']', '', regex=True)
-        df[col]=df[col].str.replace('[', '', regex=True)
-        df[col]=df[col].str.replace('{', '', regex=True)
-        df[col]=df[col].str.replace('}', '', regex=True)
-        df[col]=df[col].str.replace('|', '', regex=True)
+        df[col] = df[col].str.replace("'", "", regex=True)
+        df[col] = df[col].str.replace('"', "", regex=True)
+        df[col] = df[col].str.replace("]", "", regex=True)
+        df[col] = df[col].str.replace("[", "", regex=True)
+        df[col] = df[col].str.replace("{", "", regex=True)
+        df[col] = df[col].str.replace("}", "", regex=True)
+        df[col] = df[col].str.replace("|", "", regex=True)
     return df
 
-def create_dataframe(cells_pytess_result:list, max_cols:int, max_rows:int,csv_path):
-    '''Create dataframe using list of cell values of the table, also checks for valid header of dataframe
+
+def create_dataframe(cells_pytess_result: list, max_cols: int, max_rows: int, csv_path):
+    """Create dataframe using list of cell values of the table, also checks for valid header of dataframe
     Args:
         cells_pytess_result: list of strings, each element representing a cell in a table
         max_cols, max_rows: number of columns and rows
     Returns:
-        dataframe : final dataframe after all pre-processing 
-    '''
+        dataframe : final dataframe after all pre-processing
+    """
 
     headers = cells_pytess_result[:max_cols]
-    new_headers = uniquify(headers, (f' {x!s}' for x in string.ascii_lowercase))
+    new_headers = uniquify(headers, (f" {x!s}" for x in string.ascii_lowercase))
     counter = 0
 
     cells_list = cells_pytess_result[max_cols:]
@@ -274,10 +352,10 @@ def create_dataframe(cells_pytess_result:list, max_cols:int, max_rows:int,csv_pa
             df.iat[nrows, ncols] = str(cells_list[cell_idx])
             cell_idx += 1
 
-    ## To check if there are duplicate headers if result of uniquify+col == col 
+    ## To check if there are duplicate headers if result of uniquify+col == col
     ## This check removes headers when all headers are empty or if median of header word count is less than 6
     for x, col in zip(string.ascii_lowercase, new_headers):
-        if f' {x!s}' == col:
+        if f" {x!s}" == col:
             counter += 1
     header_char_count = [len(col) for col in new_headers]
 
@@ -291,42 +369,93 @@ def create_dataframe(cells_pytess_result:list, max_cols:int, max_rows:int,csv_pa
 
     return df
 
-def process_image(image, td_threshold, tsr_threshold, padd_top, padd_left, padd_bottom, padd_right, delta_xmin, delta_ymin, delta_xmax, delta_ymax, expand_rowcol_bbox_top, expand_rowcol_bbox_bottom):
-  image = image.convert('RGB')
-  model, probas, bboxes_scaled = table_detector(image, THRESHOLD_PROBA=td_threshold)
-  # plot_results_detection(model, image, probas, bboxes_scaled,  delta_xmin, delta_ymin, delta_xmax, delta_ymax) 
-  cropped_img_list = crop_tables(image, probas, bboxes_scaled, delta_xmin, delta_ymin, delta_xmax, delta_ymax)
-
-  result = []
-  for idx, unpadded_table in enumerate(cropped_img_list):
-      table = add_padding(unpadded_table, padd_top, padd_right, padd_bottom, padd_left)
-      model, probas, bboxes_scaled = table_struct_recog(table, THRESHOLD_PROBA=tsr_threshold)
-      rows, cols = generate_structure(model, table, probas, bboxes_scaled, expand_rowcol_bbox_top, expand_rowcol_bbox_bottom)
-      rows, cols = sort_table_featuresv2(rows, cols)
-      master_row, cols = individual_table_featuresv2(table, rows, cols)
-      cells_img, max_cols, max_rows = object_to_cellsv2(master_row, cols, expand_rowcol_bbox_top, expand_rowcol_bbox_bottom, padd_left)
-      sequential_cell_img_list = []
-      for k, img_list in cells_img.items():
-          for img in img_list:
-              sequential_cell_img_list.append(pytess(img))
-
-      csv_path = '/content/sample_data/table_' + str(idx)
-      df = create_dataframe(sequential_cell_img_list, max_cols, max_rows, csv_path)
-      result.append(df)
-  res = result[0].to_json()
-  return res
+
+def process_image(
+    image,
+    td_threshold,
+    tsr_threshold,
+    padd_top,
+    padd_left,
+    padd_bottom,
+    padd_right,
+    delta_xmin,
+    delta_ymin,
+    delta_xmax,
+    delta_ymax,
+    expand_rowcol_bbox_top,
+    expand_rowcol_bbox_bottom,
+):
+    image = image.convert("RGB")
+    model, probas, bboxes_scaled = table_detector(image, THRESHOLD_PROBA=td_threshold)
+    # plot_results_detection(model, image, probas, bboxes_scaled,  delta_xmin, delta_ymin, delta_xmax, delta_ymax)
+    cropped_img_list = crop_tables(
+        image, probas, bboxes_scaled, delta_xmin, delta_ymin, delta_xmax, delta_ymax
+    )
+
+    result = []
+    for idx, unpadded_table in enumerate(cropped_img_list):
+        table = add_padding(
+            unpadded_table, padd_top, padd_right, padd_bottom, padd_left
+        )
+        model, probas, bboxes_scaled = table_struct_recog(
+            table, THRESHOLD_PROBA=tsr_threshold
+        )
+        rows, cols = generate_structure(
+            model,
+            table,
+            probas,
+            bboxes_scaled,
+            expand_rowcol_bbox_top,
+            expand_rowcol_bbox_bottom,
+        )
+        rows, cols = sort_table_featuresv2(rows, cols)
+        master_row, cols = individual_table_featuresv2(table, rows, cols)
+        cells_img, max_cols, max_rows = object_to_cellsv2(
+            master_row,
+            cols,
+            expand_rowcol_bbox_top,
+            expand_rowcol_bbox_bottom,
+            padd_left,
+        )
+        sequential_cell_img_list = []
+        for k, img_list in cells_img.items():
+            for img in img_list:
+                sequential_cell_img_list.append(pytess(img))
+
+        csv_path = "/content/sample_data/table_" + str(idx)
+        df = create_dataframe(sequential_cell_img_list, max_cols, max_rows, csv_path)
+        result.append(df)
+    res = result[0].to_json()
+    return res
 
 
 title = "Interactive demo OCR: microsoft - table-transformer-detection + tesseract"
 description = "Demo for microsoft - table-transformer-detection + tesseract"
 article = "<p style='text-align: center'></p>"
-examples =[["image_0.png"]]
-
-iface = gr.Interface(fn=process_image, 
-                     inputs=[gr.Image(type="pil"), gr.Slider(0, 1, 0.9), gr.Slider(0, 1, 0.8), gr.Slider(0, 200, 100), gr.Slider(0, 200, 100), gr.Slider(0, 200, 100), gr.Slider(0, 200, 100), gr.Number(0), gr.Number(0), gr.Number(0), gr.Number(0),gr.Number(0),gr.Number(0)], 
-                     outputs="text",
-                     title=title,
-                     description=description,
-                     article=article,
-                     examples=examples)
+examples = [["image_0.png"]]
+
+iface = gr.Interface(
+    fn=process_image,
+    inputs=[
+        gr.Image(type="pil"),
+        gr.Slider(0, 1, 0.9),
+        gr.Slider(0, 1, 0.8),
+        gr.Slider(0, 200, 100),
+        gr.Slider(0, 200, 100),
+        gr.Slider(0, 200, 100),
+        gr.Slider(0, 200, 100),
+        gr.Number(0),
+        gr.Number(0),
+        gr.Number(0),
+        gr.Number(0),
+        gr.Number(0),
+        gr.Number(0),
+    ],
+    outputs="text",
+    title=title,
+    description=description,
+    article=article,
+    examples=examples,
+)
 iface.launch(debug=True)
+