test

app.py

@@ -2,7 +2,6 @@ import asyncio
 import string
 from collections import Counter
 from itertools import count, tee
-
 import cv2
 import matplotlib.pyplot as plt
 import numpy as np
@@ -10,64 +9,54 @@ import pandas as pd
 import streamlit as st
 import torch
 from PIL import Image
-from transformers import (DetrImageProcessor,
-                          TableTransformerForObjectDetection)
+from transformers import (DetrImageProcessor, TableTransformerForObjectDetection)
 from vietocr.tool.config import Cfg
 from vietocr.tool.predictor import Predictor
 
 st.set_option('deprecation.showPyplotGlobalUse', False)
 st.set_page_config(layout='wide')
-st.title("Table Detection and Table Structure Recognition By VWITS")
+st.title("Table Detection and Table Structure Recognition")
 st.write(
     "Implemented by MSFT team: https://github.com/microsoft/table-transformer")
 
+# Config (optional, comment out if not using)
 # config = Cfg.load_config_from_name('vgg_transformer')
-config = Cfg.load_config_from_name('vgg_seq2seq')
-config['cnn']['pretrained'] = False
-config['device'] = 'cpu'
-config['predictor']['beamsearch'] = False
-detector = Predictor(config)
+# config = Cfg.load_config_from_name('vgg_seq2seq')
+# config['cnn']['pretrained'] = False
+# config['device'] = 'cpu'
+# config['predictor']['beamsearch'] = False
+# detector = Predictor(config)
 
 table_detection_model = TableTransformerForObjectDetection.from_pretrained(
     "microsoft/table-transformer-detection")
-
 table_recognition_model = TableTransformerForObjectDetection.from_pretrained(
     "microsoft/table-transformer-structure-recognition")
 
-
 def PIL_to_cv(pil_img):
     return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
 
-
 def cv_to_PIL(cv_img):
     return Image.fromarray(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB))
 
-
 async def pytess(cell_pil_img, threshold: float = 0.5):
-    text, prob = detector.predict(cell_pil_img, return_prob=True)
+    text, prob = detector.predict(cell_pil_img, return_prob=True)  # Assuming detector is defined
     if prob < threshold:
         return ""
     return text.strip()
 
-
 def sharpen_image(pil_img):
-
     img = PIL_to_cv(pil_img)
     sharpen_kernel = np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]])
-
     sharpen = cv2.filter2D(img, -1, sharpen_kernel)
     pil_img = cv_to_PIL(sharpen)
     return pil_img
 
-
 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]
-
     suff_gens = dict(zip(not_unique, tee(suffs, len(not_unique))))
     for idx, s in enumerate(seq):
         try:
@@ -76,494 +65,20 @@ def uniquify(seq, suffs=count(1)):
             continue
         else:
             seq[idx] += suffix
-
     return seq
 
-
 def binarizeBlur_image(pil_img):
     image = PIL_to_cv(pil_img)
     thresh = cv2.threshold(image, 150, 255, cv2.THRESH_BINARY_INV)[1]
-
     result = cv2.GaussianBlur(thresh, (5, 5), 0)
     result = 255 - result
     return cv_to_PIL(result)
 
-
 def td_postprocess(pil_img):
     '''
     Removes gray background from tables
     '''
     img = PIL_to_cv(pil_img)
-
     hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
-    mask = cv2.inRange(hsv, (0, 0, 100),
-                       (255, 5, 255))  # (0, 0, 100), (255, 5, 255)
-    nzmask = cv2.inRange(hsv, (0, 0, 5),
-                         (255, 255, 255))  # (0, 0, 5), (255, 255, 255))
-    nzmask = cv2.erode(nzmask, np.ones((3, 3)))  # (3,3)
-    mask = mask & nzmask
-
-    new_img = img.copy()
-    new_img[np.where(mask)] = 255
-
-    return cv_to_PIL(new_img)
-
-
-# def super_res(pil_img):
-#     # requires opencv-contrib-python installed without the opencv-python
-#     sr = dnn_superres.DnnSuperResImpl_create()
-#     image = PIL_to_cv(pil_img)
-#     model_path = "./LapSRN_x8.pb"
-#     model_name = model_path.split('/')[1].split('_')[0].lower()
-#     model_scale = int(model_path.split('/')[1].split('_')[1].split('.')[0][1])
-
-#     sr.readModel(model_path)
-#     sr.setModel(model_name, model_scale)
-#     final_img = sr.upsample(image)
-#     final_img = cv_to_PIL(final_img)
-
-#     return final_img
-
-
-def table_detector(image, THRESHOLD_PROBA):
-    '''
-    Table detection using DEtect-object TRansformer pre-trained on 1 million tables
-
-    '''
-
-    feature_extractor = DetrImageProcessor(do_resize=True,
-                                           size=800,
-                                           max_size=800)
-    encoding = feature_extractor(image, return_tensors="pt")
-
-    with torch.no_grad():
-        outputs = table_detection_model(**encoding)
-
-    probas = outputs.logits.softmax(-1)[0, :, :-1]
-    keep = probas.max(-1).values > 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]
-
-    return (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 = DetrImageProcessor(do_resize=True,
-                                           size=1000,
-                                           max_size=1000)
-    encoding = feature_extractor(image, return_tensors="pt")
-
-    with torch.no_grad():
-        outputs = table_recognition_model(**encoding)
-
-    probas = outputs.logits.softmax(-1)[0, :, :-1]
-    keep = probas.max(-1).values > 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]
-
-    return (probas[keep], bboxes_scaled)
-
-
-class TableExtractionPipeline():
-
-    colors = ["red", "blue", "green", "yellow", "orange", "violet"]
-
-    # colors = ["red", "blue", "green", "red", "red", "red"]
-
-    def add_padding(self,
-                    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
-        result = Image.new(pil_img.mode, (new_width, new_height), color)
-        result.paste(pil_img, (left, top))
-        return result
-
-    def plot_results_detection(self, c1, model, 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 
-        '''
-        # st.write('img_obj')
-        # st.write(pil_img)
-        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 - 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')
-        c1.pyplot()
-
-    def crop_tables(self, 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 
-        '''
-        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
-            cropped_img = pil_img.crop((xmin, ymin, xmax, ymax))
-            cropped_img_list.append(cropped_img)
-
-        return cropped_img_list
-
-    def generate_structure(self, c2, model, pil_img, prob, boxes,
-                           expand_rowcol_bbox_top, expand_rowcol_bbox_bottom):
-        '''
-        Co-ordinates are adjusted here by 3 'pixels'
-        To plot table pillow image and the TSR bounding boxes on the table
-        '''
-        # st.write('img_obj')
-        # st.write(pil_img)
-        plt.figure(figsize=(32, 20))
-        plt.imshow(pil_img)
-        ax = plt.gca()
-        rows = {}
-        cols = {}
-        idx = 0
-
-        for p, (xmin, ymin, xmax, ymax) in zip(prob, boxes.tolist()):
-
-            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}'
-            # 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=self.colors[cl.item()],
-                                  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')
-        c2.pyplot()
-        return rows, cols
-
-    def sort_table_featuresv2(self, 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])
-        }
-
-        return rows_, cols_
-
-    def individual_table_featuresv2(self, pil_img, rows: dict, cols: dict):
-
-        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 object_to_cellsv2(self, 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
-        previous_xmax_col = 0
-        new_cols = {}
-        new_master_row = {}
-        previous_ymin_row = 0
-        new_cols = cols
-        new_master_row = master_row
-        ## Below 2 for loops remove redundant bounding boxes ###
-        # for k_col, v_col in cols.items():
-        #     xmin_col, _, xmax_col, _, col_img = v_col
-        #     if (np.isclose(previous_xmax_col, xmax_col, atol=5)) or (xmin_col >= xmax_col):
-        #         print('Found a column with double bbox')
-        #         continue
-        #     previous_xmax_col = xmax_col
-        #     new_cols[k_col] = v_col
-
-        # for k_row, v_row in master_row.items():
-        #     _, ymin_row, _, ymax_row, row_img = v_row
-        #     if (np.isclose(previous_ymin_row, ymin_row, atol=5)) or (ymin_row >= ymax_row):
-        #         print('Found a row with double bbox')
-        #         continue
-        #     previous_ymin_row = ymin_row
-        #     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
-            row_img_list = []
-            # plt.imshow(row_img)
-            # st.pyplot()
-            for idx, kv in enumerate(new_cols.items()):
-                k_col, v_col = kv
-                xmin_col, _, xmax_col, _, col_img = v_col
-                xmin_col, xmax_col = xmin_col - padd_left - 10, xmax_col - padd_left
-                xa = xmin_col
-                xb = xmax_col
-                if idx == 0:
-                    xa = 0
-                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
-            row_idx += 1
-
-        return cells_img, len(new_cols), len(new_master_row) - 1
-
-    def clean_dataframe(self, 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)
-        return df
-
-    @st.cache
-    def convert_df(self, df):
-        return df.to_csv().encode('utf-8')
-
-    def create_dataframe(self, c3, cell_ocr_res: list, max_cols: int,
-                         max_rows: int):
-        '''Create dataframe using list of cell values of the table, also checks for valid header of dataframe
-        Args:
-            cell_ocr_res: 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 
-        '''
-
-        headers = cell_ocr_res[:max_cols]
-        new_headers = uniquify(headers,
-                               (f' {x!s}' for x in string.ascii_lowercase))
-        counter = 0
-
-        cells_list = cell_ocr_res[max_cols:]
-        df = pd.DataFrame("", index=range(0, max_rows), columns=new_headers)
-
-        cell_idx = 0
-        for nrows in range(max_rows):
-            for ncols in range(max_cols):
-                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
-        ## 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:
-                counter += 1
-        header_char_count = [len(col) for col in new_headers]
-
-        # if (counter == len(new_headers)) or (statistics.median(header_char_count) < 6):
-        #     st.write('woooot')
-        #     df.columns = uniquify(df.iloc[0], (f' {x!s}' for x in string.ascii_lowercase))
-        #     df = df.iloc[1:,:]
-
-        df = self.clean_dataframe(df)
-
-        c3.dataframe(df)
-        csv = self.convert_df(df)
-        c3.download_button("Download table",
-                           csv,
-                           "file.csv",
-                           "text/csv",
-                           key='download-csv')
-
-        return df
-
-    async def start_process(self, image_path: str, TD_THRESHOLD, TSR_THRESHOLD,
-                            OCR_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):
-        '''
-        Initiates process of generating pandas dataframes from raw pdf-page images
-
-        '''
-        image = Image.open(image_path).convert("RGB")
-        probas, bboxes_scaled = table_detector(image,
-                                               THRESHOLD_PROBA=TD_THRESHOLD)
-
-        if bboxes_scaled.nelement() == 0:
-            st.write('No table found in the pdf-page image')
-            return ''
-
-        # try:
-        # st.write('Document: '+image_path.split('/')[-1])
-        c1, c2, c3 = st.columns((1, 1, 1))
-
-        self.plot_results_detection(c1, table_detection_model, image, probas,
-                                    bboxes_scaled, delta_xmin, delta_ymin,
-                                    delta_xmax, delta_ymax)
-        cropped_img_list = self.crop_tables(image, probas, bboxes_scaled,
-                                            delta_xmin, delta_ymin, delta_xmax,
-                                            delta_ymax)
-
-        for unpadded_table in cropped_img_list:
-
-            table = self.add_padding(unpadded_table, padd_top, padd_right,
-                                     padd_bottom, padd_left)
-            # table = super_res(table)
-            # table = binarizeBlur_image(table)
-            # table = sharpen_image(table) # Test sharpen image next
-            # table = td_postprocess(table)
-
-            probas, bboxes_scaled = table_struct_recog(
-                table, THRESHOLD_PROBA=TSR_THRESHOLD)
-            rows, cols = self.generate_structure(c2, table_recognition_model,
-                                                 table, probas, bboxes_scaled,
-                                                 expand_rowcol_bbox_top,
-                                                 expand_rowcol_bbox_bottom)
-            # st.write(len(rows), len(cols))
-            rows, cols = self.sort_table_featuresv2(rows, cols)
-            master_row, cols = self.individual_table_featuresv2(
-                table, rows, cols)
-
-            cells_img, max_cols, max_rows = self.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:
-                    # img = super_res(img)
-                    # img = sharpen_image(img) # Test sharpen image next
-                    # img = binarizeBlur_image(img)
-                    # img = self.add_padding(img, 10,10,10,10)
-                    # plt.imshow(img)
-                    # c3.pyplot()
-                    sequential_cell_img_list.append(
-                        pytess(cell_pil_img=img, threshold=OCR_THRESHOLD))
-
-            cell_ocr_res = await asyncio.gather(*sequential_cell_img_list)
-
-            self.create_dataframe(c3, cell_ocr_res, max_cols, max_rows)
-            st.write(
-                'Errors in OCR is due to either quality of the image or performance of the OCR'
-            )
-        # except:
-        #     st.write('Either incorrectly identified table or no table, to debug remove try/except')
-        # break
-        # break
-
-
-if __name__ == "__main__":
-
-    img_name = st.file_uploader("Upload an image with table(s)")
-    st1, st2, st3 = st.columns((1, 1, 1))
-    TD_th = st1.slider('Table detection threshold', 0.0, 1.0, 0.8)
-    TSR_th = st2.slider('Table structure recognition threshold', 0.0, 1.0, 0.8)
-    OCR_th = st3.slider("Text Probs Threshold", 0.0, 1.0, 0.5)
-
-    st1, st2, st3, st4 = st.columns((1, 1, 1, 1))
-
-    padd_top = st1.slider('Padding top', 0, 200, 40)
-    padd_left = st2.slider('Padding left', 0, 200, 40)
-    padd_right = st3.slider('Padding right', 0, 200, 40)
-    padd_bottom = st4.slider('Padding bottom', 0, 200, 40)
-
-    te = TableExtractionPipeline()
-    # for img in image_list:
-    if img_name is not None:
-        asyncio.run(
-            te.start_process(img_name,
-                             TD_THRESHOLD=TD_th,
-                             TSR_THRESHOLD=TSR_th,
-                             OCR_THRESHOLD=OCR_th,
-                             padd_top=padd_top,
-                             padd_left=padd_left,
-                             padd_bottom=padd_bottom,
-                             padd_right=padd_right,
-                             delta_xmin=0,
-                             delta_ymin=0,
-                             delta_xmax=0,
-                             delta_ymax=0,
-                             expand_rowcol_bbox_top=0,
-                             expand_rowcol_bbox_bottom=0))
+    mask = cv2.inRange(hsv, (0, 0, 100), (255, 5, 255))  # (0, 0, 100), (255, 5, 255)
+    nzmask = cv2.inRange(hsv, (0, 0, 5), (255, 255, 255))  # (0, 0,