trysem/TableIMG2-CSV
1
1import streamlit as st2from PIL import Image, ImageEnhance3import statistics4import os5import string6from collections import Counter 7from itertools import tee, count8# import TDTSR9import pytesseract10from pytesseract import Output11import json12import pandas as pd13import matplotlib.pyplot as plt14import cv215import numpy as np16# from transformers import TrOCRProcessor, VisionEncoderDecoderModel17# from cv2 import dnn_superres18from transformers import DetrFeatureExtractor19#from transformers import DetrForObjectDetection20from transformers import TableTransformerForObjectDetection21import torch22import asyncio23# pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'24 25 26st.set_option('deprecation.showPyplotGlobalUse', False)27st.set_page_config(layout='wide')28st.title("Table Detection and Table Structure Recognition")29st.write("Implemented by MSFT team: https://github.com/microsoft/table-transformer")30 31 32 33def PIL_to_cv(pil_img):34 return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR) 35 36def cv_to_PIL(cv_img):37 return Image.fromarray(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB))38 39 40async def pytess(cell_pil_img):41 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()42 43 44# def super_res(pil_img):45 # '''46 # Useful for low-res docs47 # '''48 # requires opencv-contrib-python installed without the opencv-python49 # sr = dnn_superres.DnnSuperResImpl_create()50 # image = PIL_to_cv(pil_img)51 # model_path = "/data/Salman/TRD/code/table-transformer/transformers/LapSRN_x2.pb"52 # model_name = 'lapsrn'53 # model_scale = 254 # sr.readModel(model_path)55 # sr.setModel(model_name, model_scale)56 # final_img = sr.upsample(image)57 # final_img = cv_to_PIL(final_img)58 59 # return final_img60 61 62def sharpen_image(pil_img):63 64 img = PIL_to_cv(pil_img)65 sharpen_kernel = np.array([[-1, -1, -1], 66 [-1, 9, -1], 67 [-1, -1, -1]])68 69 sharpen = cv2.filter2D(img, -1, sharpen_kernel)70 pil_img = cv_to_PIL(sharpen)71 return pil_img72 73 74def uniquify(seq, suffs = count(1)):75 """Make all the items unique by adding a suffix (1, 2, etc).76 Credit: https://stackoverflow.com/questions/30650474/python-rename-duplicates-in-list-with-progressive-numbers-without-sorting-list77 `seq` is mutable sequence of strings.78 `suffs` is an optional alternative suffix iterable.79 """80 not_unique = [k for k,v in Counter(seq).items() if v>1] 81 82 suff_gens = dict(zip(not_unique, tee(suffs, len(not_unique)))) 83 for idx,s in enumerate(seq):84 try:85 suffix = str(next(suff_gens[s]))86 except KeyError:87 continue88 else:89 seq[idx] += suffix90 91 return seq92 93def binarizeBlur_image(pil_img):94 image = PIL_to_cv(pil_img)95 thresh = cv2.threshold(image, 150, 255, cv2.THRESH_BINARY_INV)[1]96 97 result = cv2.GaussianBlur(thresh, (5,5), 0)98 result = 255 - result99 return cv_to_PIL(result)100 101 102 103def td_postprocess(pil_img):104 '''105 Removes gray background from tables106 '''107 img = PIL_to_cv(pil_img)108 109 hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)110 mask = cv2.inRange(hsv, (0, 0, 100), (255, 5, 255)) # (0, 0, 100), (255, 5, 255)111 nzmask = cv2.inRange(hsv, (0, 0, 5), (255, 255, 255)) # (0, 0, 5), (255, 255, 255))112 nzmask = cv2.erode(nzmask, np.ones((3,3))) # (3,3)113 mask = mask & nzmask114 115 new_img = img.copy()116 new_img[np.where(mask)] = 255117 118 119 return cv_to_PIL(new_img)120 121# def super_res(pil_img):122# # requires opencv-contrib-python installed without the opencv-python123# sr = dnn_superres.DnnSuperResImpl_create()124# image = PIL_to_cv(pil_img)125# model_path = "./LapSRN_x8.pb"126# model_name = model_path.split('/')[1].split('_')[0].lower()127# model_scale = int(model_path.split('/')[1].split('_')[1].split('.')[0][1])128 129# sr.readModel(model_path)130# sr.setModel(model_name, model_scale)131# final_img = sr.upsample(image)132# final_img = cv_to_PIL(final_img)133 134# return final_img135 136def table_detector(image, THRESHOLD_PROBA):137 '''138 Table detection using DEtect-object TRansformer pre-trained on 1 million tables139 140 '''141 142 feature_extractor = DetrFeatureExtractor(do_resize=True, size=800, max_size=800)143 encoding = feature_extractor(image, return_tensors="pt")144 145 model = TableTransformerForObjectDetection.from_pretrained("microsoft/table-transformer-detection")146 147 with torch.no_grad():148 outputs = model(**encoding)149 150 probas = outputs.logits.softmax(-1)[0, :, :-1]151 keep = probas.max(-1).values > THRESHOLD_PROBA152 153 target_sizes = torch.tensor(image.size[::-1]).unsqueeze(0)154 postprocessed_outputs = feature_extractor.post_process(outputs, target_sizes)155 bboxes_scaled = postprocessed_outputs[0]['boxes'][keep]156 157 return (model, probas[keep], bboxes_scaled)158 159 160def table_struct_recog(image, THRESHOLD_PROBA):161 '''162 Table structure recognition using DEtect-object TRansformer pre-trained on 1 million tables163 '''164 165 feature_extractor = DetrFeatureExtractor(do_resize=True, size=1000, max_size=1000)166 encoding = feature_extractor(image, return_tensors="pt")167 168 model = TableTransformerForObjectDetection.from_pretrained("microsoft/table-transformer-structure-recognition")169 with torch.no_grad():170 outputs = model(**encoding)171 172 probas = outputs.logits.softmax(-1)[0, :, :-1]173 keep = probas.max(-1).values > THRESHOLD_PROBA174 175 target_sizes = torch.tensor(image.size[::-1]).unsqueeze(0)176 postprocessed_outputs = feature_extractor.post_process(outputs, target_sizes)177 bboxes_scaled = postprocessed_outputs[0]['boxes'][keep]178 179 return (model, probas[keep], bboxes_scaled)180 181 182 183 184 185class TableExtractionPipeline():186 187 colors = ["red", "blue", "green", "yellow", "orange", "violet"]188 189 # colors = ["red", "blue", "green", "red", "red", "red"]190 191 def add_padding(self, pil_img, top, right, bottom, left, color=(255,255,255)):192 '''193 Image padding as part of TSR pre-processing to prevent missing table edges194 '''195 width, height = pil_img.size196 new_width = width + right + left197 new_height = height + top + bottom198 result = Image.new(pil_img.mode, (new_width, new_height), color)199 result.paste(pil_img, (left, top))200 return result201 202 def plot_results_detection(self, c1, model, pil_img, prob, boxes, delta_xmin, delta_ymin, delta_xmax, delta_ymax):203 '''204 crop_tables and plot_results_detection must have same co-ord shifts because 1 only plots the other one updates co-ordinates 205 '''206 # st.write('img_obj')207 # st.write(pil_img)208 plt.imshow(pil_img)209 ax = plt.gca()210 211 for p, (xmin, ymin, xmax, ymax) in zip(prob, boxes.tolist()):212 cl = p.argmax()213 xmin, ymin, xmax, ymax = xmin-delta_xmin, ymin-delta_ymin, xmax+delta_xmax, ymax+delta_ymax 214 ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,fill=False, color='red', linewidth=3))215 text = f'{model.config.id2label[cl.item()]}: {p[cl]:0.2f}'216 ax.text(xmin-20, ymin-50, text, fontsize=10,bbox=dict(facecolor='yellow', alpha=0.5))217 plt.axis('off')218 c1.pyplot()219 220 221 def crop_tables(self, pil_img, prob, boxes, delta_xmin, delta_ymin, delta_xmax, delta_ymax):222 '''223 crop_tables and plot_results_detection must have same co-ord shifts because 1 only plots the other one updates co-ordinates 224 '''225 cropped_img_list = []226 227 for p, (xmin, ymin, xmax, ymax) in zip(prob, boxes.tolist()):228 229 xmin, ymin, xmax, ymax = xmin-delta_xmin, ymin-delta_ymin, xmax+delta_xmax, ymax+delta_ymax 230 cropped_img = pil_img.crop((xmin, ymin, xmax, ymax))231 cropped_img_list.append(cropped_img)232 233 234 return cropped_img_list235 236 def generate_structure(self, c2, model, pil_img, prob, boxes, expand_rowcol_bbox_top, expand_rowcol_bbox_bottom):237 '''238 Co-ordinates are adjusted here by 3 'pixels'239 To plot table pillow image and the TSR bounding boxes on the table240 '''241 # st.write('img_obj')242 # st.write(pil_img)243 plt.figure(figsize=(32,20))244 plt.imshow(pil_img)245 ax = plt.gca()246 rows = {}247 cols = {}248 idx = 0249 250 251 for p, (xmin, ymin, xmax, ymax) in zip(prob, boxes.tolist()):252 253 xmin, ymin, xmax, ymax = xmin, ymin, xmax, ymax 254 cl = p.argmax()255 class_text = model.config.id2label[cl.item()]256 text = f'{class_text}: {p[cl]:0.2f}'257 # or (class_text == 'table column')258 if (class_text == 'table row') or (class_text =='table projected row header') or (class_text == 'table column'):259 ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,fill=False, color=self.colors[cl.item()], linewidth=2))260 ax.text(xmin-10, ymin-10, text, fontsize=5, bbox=dict(facecolor='yellow', alpha=0.5))261 262 if class_text == 'table row':263 rows['table row.'+str(idx)] = (xmin, ymin-expand_rowcol_bbox_top, xmax, ymax+expand_rowcol_bbox_bottom)264 if class_text == 'table column':265 cols['table column.'+str(idx)] = (xmin, ymin-expand_rowcol_bbox_top, xmax, ymax+expand_rowcol_bbox_bottom)266 267 idx += 1268 269 270 plt.axis('on')271 c2.pyplot()272 return rows, cols273 274 def sort_table_featuresv2(self, rows:dict, cols:dict):275 # Sometimes the header and first row overlap, and we need the header bbox not to have first row's bbox inside the headers bbox276 rows_ = {table_feature : (xmin, ymin, xmax, ymax) for table_feature, (xmin, ymin, xmax, ymax) in sorted(rows.items(), key=lambda tup: tup[1][1])}277 cols_ = {table_feature : (xmin, ymin, xmax, ymax) for table_feature, (xmin, ymin, xmax, ymax) in sorted(cols.items(), key=lambda tup: tup[1][0])}278 279 return rows_, cols_280 281 def individual_table_featuresv2(self, pil_img, rows:dict, cols:dict):282 283 for k, v in rows.items():284 xmin, ymin, xmax, ymax = v285 cropped_img = pil_img.crop((xmin, ymin, xmax, ymax))286 rows[k] = xmin, ymin, xmax, ymax, cropped_img287 288 for k, v in cols.items():289 xmin, ymin, xmax, ymax = v290 cropped_img = pil_img.crop((xmin, ymin, xmax, ymax))291 cols[k] = xmin, ymin, xmax, ymax, cropped_img292 293 return rows, cols294 295 296 def object_to_cellsv2(self, master_row:dict, cols:dict, expand_rowcol_bbox_top, expand_rowcol_bbox_bottom, padd_left):297 '''Removes redundant bbox for rows&columns and divides each row into cells from columns298 Args:299 300 Returns:301 302 303 '''304 cells_img = {}305 header_idx = 0306 row_idx = 0307 previous_xmax_col = 0308 new_cols = {}309 new_master_row = {}310 previous_ymin_row = 0311 new_cols = cols312 new_master_row = master_row313 ## Below 2 for loops remove redundant bounding boxes ###314 # for k_col, v_col in cols.items():315 # xmin_col, _, xmax_col, _, col_img = v_col316 # if (np.isclose(previous_xmax_col, xmax_col, atol=5)) or (xmin_col >= xmax_col):317 # print('Found a column with double bbox')318 # continue319 # previous_xmax_col = xmax_col320 # new_cols[k_col] = v_col321 322 # for k_row, v_row in master_row.items():323 # _, ymin_row, _, ymax_row, row_img = v_row324 # if (np.isclose(previous_ymin_row, ymin_row, atol=5)) or (ymin_row >= ymax_row):325 # print('Found a row with double bbox')326 # continue327 # previous_ymin_row = ymin_row328 # new_master_row[k_row] = v_row329 ######################################################330 for k_row, v_row in new_master_row.items():331 332 _, _, _, _, row_img = v_row333 xmax, ymax = row_img.size334 xa, ya, xb, yb = 0, 0, 0, ymax335 row_img_list = []336 # plt.imshow(row_img)337 # st.pyplot()338 for idx, kv in enumerate(new_cols.items()):339 k_col, v_col = kv340 xmin_col, _, xmax_col, _, col_img = v_col341 xmin_col, xmax_col = xmin_col - padd_left - 10, xmax_col - padd_left342 # plt.imshow(col_img)343 # st.pyplot()344 # xa + 3 : to remove borders on the left side of the cropped cell345 # yb = 3: to remove row information from the above row of the cropped cell346 # xb - 3: to remove borders on the right side of the cropped cell347 xa = xmin_col348 xb = xmax_col349 if idx == 0:350 xa = 0351 if idx == len(new_cols)-1:352 xb = xmax353 xa, ya, xb, yb = xa, ya, xb, yb354 355 row_img_cropped = row_img.crop((xa, ya, xb, yb))356 row_img_list.append(row_img_cropped)357 358 cells_img[k_row+'.'+str(row_idx)] = row_img_list359 row_idx += 1360 361 return cells_img, len(new_cols), len(new_master_row)-1362 363 def clean_dataframe(self, df):364 '''365 Remove irrelevant symbols that appear with tesseractOCR366 '''367 # df.columns = [col.replace('|', '') for col in df.columns]368 369 for col in df.columns:370 371 df[col]=df[col].str.replace("'", '', regex=True)372 df[col]=df[col].str.replace('"', '', regex=True)373 df[col]=df[col].str.replace(']', '', regex=True)374 df[col]=df[col].str.replace('[', '', regex=True)375 df[col]=df[col].str.replace('{', '', regex=True)376 df[col]=df[col].str.replace('}', '', regex=True)377 return df378 379 @st.cache380 def convert_df(self, df):381 return df.to_csv().encode('utf-8')382 383 384 def create_dataframe(self, c3, cells_pytess_result:list, max_cols:int, max_rows:int):385 '''Create dataframe using list of cell values of the table, also checks for valid header of dataframe386 Args:387 cells_pytess_result: list of strings, each element representing a cell in a table388 max_cols, max_rows: number of columns and rows389 Returns:390 dataframe : final dataframe after all pre-processing 391 '''392 393 headers = cells_pytess_result[:max_cols]394 new_headers = uniquify(headers, (f' {x!s}' for x in string.ascii_lowercase))395 counter = 0396 397 cells_list = cells_pytess_result[max_cols:]398 df = pd.DataFrame("", index=range(0, max_rows), columns=new_headers)399 400 cell_idx = 0401 for nrows in range(max_rows):402 for ncols in range(max_cols):403 df.iat[nrows, ncols] = str(cells_list[cell_idx])404 cell_idx += 1405 406 ## To check if there are duplicate headers if result of uniquify+col == col 407 ## This check removes headers when all headers are empty or if median of header word count is less than 6408 for x, col in zip(string.ascii_lowercase, new_headers):409 if f' {x!s}' == col:410 counter += 1411 header_char_count = [len(col) for col in new_headers]412 413 # if (counter == len(new_headers)) or (statistics.median(header_char_count) < 6):414 # st.write('woooot')415 # df.columns = uniquify(df.iloc[0], (f' {x!s}' for x in string.ascii_lowercase))416 # df = df.iloc[1:,:]417 418 df = self.clean_dataframe(df)419 420 c3.dataframe(df)421 csv = self.convert_df(df)422 c3.download_button("Download table", csv, "file.csv", "text/csv", key='download-csv')423 424 return df425 426 427 428 429 430 431 async def start_process(self, image_path:str, 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):432 '''433 Initiates process of generating pandas dataframes from raw pdf-page images434 435 '''436 image = Image.open(image_path).convert("RGB")437 model, probas, bboxes_scaled = table_detector(image, THRESHOLD_PROBA=TD_THRESHOLD)438 439 if bboxes_scaled.nelement() == 0:440 st.write('No table found in the pdf-page image')441 return ''442 443 # try:444 # st.write('Document: '+image_path.split('/')[-1])445 c1, c2, c3 = st.columns((1,1,1))446 447 self.plot_results_detection(c1, model, image, probas, bboxes_scaled, delta_xmin, delta_ymin, delta_xmax, delta_ymax) 448 cropped_img_list = self.crop_tables(image, probas, bboxes_scaled, delta_xmin, delta_ymin, delta_xmax, delta_ymax)449 450 for unpadded_table in cropped_img_list:451 452 table = self.add_padding(unpadded_table, padd_top, padd_right, padd_bottom, padd_left)453 # table = super_res(table)454 # table = binarizeBlur_image(table)455 # table = sharpen_image(table) # Test sharpen image next456 # table = td_postprocess(table)457 458 model, probas, bboxes_scaled = table_struct_recog(table, THRESHOLD_PROBA=TSR_THRESHOLD)459 rows, cols = self.generate_structure(c2, model, table, probas, bboxes_scaled, expand_rowcol_bbox_top, expand_rowcol_bbox_bottom)460 # st.write(len(rows), len(cols))461 rows, cols = self.sort_table_featuresv2(rows, cols)462 master_row, cols = self.individual_table_featuresv2(table, rows, cols)463 464 cells_img, max_cols, max_rows = self.object_to_cellsv2(master_row, cols, expand_rowcol_bbox_top, expand_rowcol_bbox_bottom, padd_left)465 466 sequential_cell_img_list = []467 for k, img_list in cells_img.items():468 for img in img_list:469 # img = super_res(img)470 # img = sharpen_image(img) # Test sharpen image next471 # img = binarizeBlur_image(img)472 # img = self.add_padding(img, 10,10,10,10)473 # plt.imshow(img)474 # c3.pyplot()475 sequential_cell_img_list.append(pytess(img))476 477 cells_pytess_result = await asyncio.gather(*sequential_cell_img_list)478 479 480 self.create_dataframe(c3, cells_pytess_result, max_cols, max_rows)481 st.write('Errors in OCR is due to either quality of the image or performance of the OCR')482 # except:483 # st.write('Either incorrectly identified table or no table, to debug remove try/except')484 # break485 # break486 487 488 489 490if __name__ == "__main__":491 492 img_name = st.file_uploader("Upload an image with table(s)")493 st1, st2 = st.columns((1,1))494 TD_th = st1.slider('Table detection threshold', 0.0, 1.0, 0.6)495 TSR_th = st2.slider('Table structure recognition threshold', 0.0, 1.0, 0.8)496 497 st1, st2, st3, st4 = st.columns((1,1,1,1))498 499 padd_top = st1.slider('Padding top', 0, 200, 20)500 padd_left = st2.slider('Padding left', 0, 200, 20)501 padd_right = st3.slider('Padding right', 0, 200, 20)502 padd_bottom = st4.slider('Padding bottom', 0, 200, 20)503 504 te = TableExtractionPipeline()505 # for img in image_list:506 if img_name is not None:507 asyncio.run(te.start_process(img_name, TD_THRESHOLD=TD_th , TSR_THRESHOLD=TSR_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))508 509 510 511