Svngoku SalML commited on
Commit
d88d2c7
0 Parent(s):

Duplicate from SalML/TableTransformer2CSV

Browse files

Co-authored-by: Mohammed Salman <SalML@users.noreply.huggingface.co>

Files changed (5) hide show
  1. .gitattributes +31 -0
  2. README.md +13 -0
  3. app.py +509 -0
  4. packages.txt +6 -0
  5. requirements.txt +47 -0
.gitattributes ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ *.7z filter=lfs diff=lfs merge=lfs -text
2
+ *.arrow filter=lfs diff=lfs merge=lfs -text
3
+ *.bin filter=lfs diff=lfs merge=lfs -text
4
+ *.bz2 filter=lfs diff=lfs merge=lfs -text
5
+ *.ftz filter=lfs diff=lfs merge=lfs -text
6
+ *.gz filter=lfs diff=lfs merge=lfs -text
7
+ *.h5 filter=lfs diff=lfs merge=lfs -text
8
+ *.joblib filter=lfs diff=lfs merge=lfs -text
9
+ *.lfs.* filter=lfs diff=lfs merge=lfs -text
10
+ *.model filter=lfs diff=lfs merge=lfs -text
11
+ *.msgpack filter=lfs diff=lfs merge=lfs -text
12
+ *.npy filter=lfs diff=lfs merge=lfs -text
13
+ *.npz filter=lfs diff=lfs merge=lfs -text
14
+ *.onnx filter=lfs diff=lfs merge=lfs -text
15
+ *.ot filter=lfs diff=lfs merge=lfs -text
16
+ *.parquet filter=lfs diff=lfs merge=lfs -text
17
+ *.pickle filter=lfs diff=lfs merge=lfs -text
18
+ *.pkl filter=lfs diff=lfs merge=lfs -text
19
+ *.pb filter=lfs diff=lfs merge=lfs -text
20
+ *.pt filter=lfs diff=lfs merge=lfs -text
21
+ *.pth filter=lfs diff=lfs merge=lfs -text
22
+ *.rar filter=lfs diff=lfs merge=lfs -text
23
+ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
24
+ *.tar.* filter=lfs diff=lfs merge=lfs -text
25
+ *.tflite filter=lfs diff=lfs merge=lfs -text
26
+ *.tgz filter=lfs diff=lfs merge=lfs -text
27
+ *.wasm filter=lfs diff=lfs merge=lfs -text
28
+ *.xz filter=lfs diff=lfs merge=lfs -text
29
+ *.zip filter=lfs diff=lfs merge=lfs -text
30
+ *.zst filter=lfs diff=lfs merge=lfs -text
31
+ *tfevents* filter=lfs diff=lfs merge=lfs -text
README.md ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ title: Image2Table
3
+ emoji: 🚀
4
+ colorFrom: indigo
5
+ colorTo: purple
6
+ sdk: streamlit
7
+ sdk_version: 1.10.0
8
+ app_file: app.py
9
+ pinned: false
10
+ duplicated_from: SalML/TableTransformer2CSV
11
+ ---
12
+
13
+ Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
app.py ADDED
@@ -0,0 +1,509 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from PIL import Image, ImageEnhance
3
+ import statistics
4
+ import os
5
+ import string
6
+ from collections import Counter
7
+ from itertools import tee, count
8
+ # import TDTSR
9
+ import pytesseract
10
+ from pytesseract import Output
11
+ import json
12
+ import pandas as pd
13
+ import matplotlib.pyplot as plt
14
+ import cv2
15
+ import numpy as np
16
+ # from transformers import TrOCRProcessor, VisionEncoderDecoderModel
17
+ # from cv2 import dnn_superres
18
+ from transformers import DetrFeatureExtractor
19
+ #from transformers import DetrForObjectDetection
20
+ from transformers import TableTransformerForObjectDetection
21
+ import torch
22
+ import asyncio
23
+ # pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'
24
+
25
+
26
+ st.set_option('deprecation.showPyplotGlobalUse', False)
27
+ st.set_page_config(layout='wide')
28
+ st.title("Table Detection and Table Structure Recognition")
29
+ st.write("Implemented by MSFT team: https://github.com/microsoft/table-transformer")
30
+
31
+
32
+ def PIL_to_cv(pil_img):
33
+ return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
34
+
35
+ def cv_to_PIL(cv_img):
36
+ return Image.fromarray(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB))
37
+
38
+
39
+ async def pytess(cell_pil_img):
40
+ 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()
41
+
42
+
43
+ # def super_res(pil_img):
44
+ # '''
45
+ # Useful for low-res docs
46
+ # '''
47
+ # requires opencv-contrib-python installed without the opencv-python
48
+ # sr = dnn_superres.DnnSuperResImpl_create()
49
+ # image = PIL_to_cv(pil_img)
50
+ # model_path = "/data/Salman/TRD/code/table-transformer/transformers/LapSRN_x2.pb"
51
+ # model_name = 'lapsrn'
52
+ # model_scale = 2
53
+ # sr.readModel(model_path)
54
+ # sr.setModel(model_name, model_scale)
55
+ # final_img = sr.upsample(image)
56
+ # final_img = cv_to_PIL(final_img)
57
+
58
+ # return final_img
59
+
60
+
61
+ def sharpen_image(pil_img):
62
+
63
+ img = PIL_to_cv(pil_img)
64
+ sharpen_kernel = np.array([[-1, -1, -1],
65
+ [-1, 9, -1],
66
+ [-1, -1, -1]])
67
+
68
+ sharpen = cv2.filter2D(img, -1, sharpen_kernel)
69
+ pil_img = cv_to_PIL(sharpen)
70
+ return pil_img
71
+
72
+
73
+ def uniquify(seq, suffs = count(1)):
74
+ """Make all the items unique by adding a suffix (1, 2, etc).
75
+ Credit: https://stackoverflow.com/questions/30650474/python-rename-duplicates-in-list-with-progressive-numbers-without-sorting-list
76
+ `seq` is mutable sequence of strings.
77
+ `suffs` is an optional alternative suffix iterable.
78
+ """
79
+ not_unique = [k for k,v in Counter(seq).items() if v>1]
80
+
81
+ suff_gens = dict(zip(not_unique, tee(suffs, len(not_unique))))
82
+ for idx,s in enumerate(seq):
83
+ try:
84
+ suffix = str(next(suff_gens[s]))
85
+ except KeyError:
86
+ continue
87
+ else:
88
+ seq[idx] += suffix
89
+
90
+ return seq
91
+
92
+ def binarizeBlur_image(pil_img):
93
+ image = PIL_to_cv(pil_img)
94
+ thresh = cv2.threshold(image, 150, 255, cv2.THRESH_BINARY_INV)[1]
95
+
96
+ result = cv2.GaussianBlur(thresh, (5,5), 0)
97
+ result = 255 - result
98
+ return cv_to_PIL(result)
99
+
100
+
101
+
102
+ def td_postprocess(pil_img):
103
+ '''
104
+ Removes gray background from tables
105
+ '''
106
+ img = PIL_to_cv(pil_img)
107
+
108
+ hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
109
+ mask = cv2.inRange(hsv, (0, 0, 100), (255, 5, 255)) # (0, 0, 100), (255, 5, 255)
110
+ nzmask = cv2.inRange(hsv, (0, 0, 5), (255, 255, 255)) # (0, 0, 5), (255, 255, 255))
111
+ nzmask = cv2.erode(nzmask, np.ones((3,3))) # (3,3)
112
+ mask = mask & nzmask
113
+
114
+ new_img = img.copy()
115
+ new_img[np.where(mask)] = 255
116
+
117
+
118
+ return cv_to_PIL(new_img)
119
+
120
+ # def super_res(pil_img):
121
+ # # requires opencv-contrib-python installed without the opencv-python
122
+ # sr = dnn_superres.DnnSuperResImpl_create()
123
+ # image = PIL_to_cv(pil_img)
124
+ # model_path = "./LapSRN_x8.pb"
125
+ # model_name = model_path.split('/')[1].split('_')[0].lower()
126
+ # model_scale = int(model_path.split('/')[1].split('_')[1].split('.')[0][1])
127
+
128
+ # sr.readModel(model_path)
129
+ # sr.setModel(model_name, model_scale)
130
+ # final_img = sr.upsample(image)
131
+ # final_img = cv_to_PIL(final_img)
132
+
133
+ # return final_img
134
+
135
+ def table_detector(image, THRESHOLD_PROBA):
136
+ '''
137
+ Table detection using DEtect-object TRansformer pre-trained on 1 million tables
138
+
139
+ '''
140
+
141
+ feature_extractor = DetrFeatureExtractor(do_resize=True, size=800, max_size=800)
142
+ encoding = feature_extractor(image, return_tensors="pt")
143
+
144
+ model = TableTransformerForObjectDetection.from_pretrained("microsoft/table-transformer-detection")
145
+
146
+ with torch.no_grad():
147
+ outputs = model(**encoding)
148
+
149
+ probas = outputs.logits.softmax(-1)[0, :, :-1]
150
+ keep = probas.max(-1).values > THRESHOLD_PROBA
151
+
152
+ target_sizes = torch.tensor(image.size[::-1]).unsqueeze(0)
153
+ postprocessed_outputs = feature_extractor.post_process(outputs, target_sizes)
154
+ bboxes_scaled = postprocessed_outputs[0]['boxes'][keep]
155
+
156
+ return (model, probas[keep], bboxes_scaled)
157
+
158
+
159
+ def table_struct_recog(image, THRESHOLD_PROBA):
160
+ '''
161
+ Table structure recognition using DEtect-object TRansformer pre-trained on 1 million tables
162
+ '''
163
+
164
+ feature_extractor = DetrFeatureExtractor(do_resize=True, size=1000, max_size=1000)
165
+ encoding = feature_extractor(image, return_tensors="pt")
166
+
167
+ model = TableTransformerForObjectDetection.from_pretrained("microsoft/table-transformer-structure-recognition")
168
+ with torch.no_grad():
169
+ outputs = model(**encoding)
170
+
171
+ probas = outputs.logits.softmax(-1)[0, :, :-1]
172
+ keep = probas.max(-1).values > THRESHOLD_PROBA
173
+
174
+ target_sizes = torch.tensor(image.size[::-1]).unsqueeze(0)
175
+ postprocessed_outputs = feature_extractor.post_process(outputs, target_sizes)
176
+ bboxes_scaled = postprocessed_outputs[0]['boxes'][keep]
177
+
178
+ return (model, probas[keep], bboxes_scaled)
179
+
180
+
181
+
182
+
183
+
184
+ class TableExtractionPipeline():
185
+
186
+ colors = ["red", "blue", "green", "yellow", "orange", "violet"]
187
+
188
+ # colors = ["red", "blue", "green", "red", "red", "red"]
189
+
190
+ def add_padding(self, pil_img, top, right, bottom, left, color=(255,255,255)):
191
+ '''
192
+ Image padding as part of TSR pre-processing to prevent missing table edges
193
+ '''
194
+ width, height = pil_img.size
195
+ new_width = width + right + left
196
+ new_height = height + top + bottom
197
+ result = Image.new(pil_img.mode, (new_width, new_height), color)
198
+ result.paste(pil_img, (left, top))
199
+ return result
200
+
201
+ def plot_results_detection(self, c1, model, pil_img, prob, boxes, delta_xmin, delta_ymin, delta_xmax, delta_ymax):
202
+ '''
203
+ crop_tables and plot_results_detection must have same co-ord shifts because 1 only plots the other one updates co-ordinates
204
+ '''
205
+ # st.write('img_obj')
206
+ # st.write(pil_img)
207
+ plt.imshow(pil_img)
208
+ ax = plt.gca()
209
+
210
+ for p, (xmin, ymin, xmax, ymax) in zip(prob, boxes.tolist()):
211
+ cl = p.argmax()
212
+ xmin, ymin, xmax, ymax = xmin-delta_xmin, ymin-delta_ymin, xmax+delta_xmax, ymax+delta_ymax
213
+ ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,fill=False, color='red', linewidth=3))
214
+ text = f'{model.config.id2label[cl.item()]}: {p[cl]:0.2f}'
215
+ ax.text(xmin-20, ymin-50, text, fontsize=10,bbox=dict(facecolor='yellow', alpha=0.5))
216
+ plt.axis('off')
217
+ c1.pyplot()
218
+
219
+
220
+ def crop_tables(self, pil_img, prob, boxes, delta_xmin, delta_ymin, delta_xmax, delta_ymax):
221
+ '''
222
+ crop_tables and plot_results_detection must have same co-ord shifts because 1 only plots the other one updates co-ordinates
223
+ '''
224
+ cropped_img_list = []
225
+
226
+ for p, (xmin, ymin, xmax, ymax) in zip(prob, boxes.tolist()):
227
+
228
+ xmin, ymin, xmax, ymax = xmin-delta_xmin, ymin-delta_ymin, xmax+delta_xmax, ymax+delta_ymax
229
+ cropped_img = pil_img.crop((xmin, ymin, xmax, ymax))
230
+ cropped_img_list.append(cropped_img)
231
+
232
+
233
+ return cropped_img_list
234
+
235
+ def generate_structure(self, c2, model, pil_img, prob, boxes, expand_rowcol_bbox_top, expand_rowcol_bbox_bottom):
236
+ '''
237
+ Co-ordinates are adjusted here by 3 'pixels'
238
+ To plot table pillow image and the TSR bounding boxes on the table
239
+ '''
240
+ # st.write('img_obj')
241
+ # st.write(pil_img)
242
+ plt.figure(figsize=(32,20))
243
+ plt.imshow(pil_img)
244
+ ax = plt.gca()
245
+ rows = {}
246
+ cols = {}
247
+ idx = 0
248
+
249
+
250
+ for p, (xmin, ymin, xmax, ymax) in zip(prob, boxes.tolist()):
251
+
252
+ xmin, ymin, xmax, ymax = xmin, ymin, xmax, ymax
253
+ cl = p.argmax()
254
+ class_text = model.config.id2label[cl.item()]
255
+ text = f'{class_text}: {p[cl]:0.2f}'
256
+ # or (class_text == 'table column')
257
+ if (class_text == 'table row') or (class_text =='table projected row header') or (class_text == 'table column'):
258
+ ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,fill=False, color=self.colors[cl.item()], linewidth=2))
259
+ ax.text(xmin-10, ymin-10, text, fontsize=5, bbox=dict(facecolor='yellow', alpha=0.5))
260
+
261
+ if class_text == 'table row':
262
+ rows['table row.'+str(idx)] = (xmin, ymin-expand_rowcol_bbox_top, xmax, ymax+expand_rowcol_bbox_bottom)
263
+ if class_text == 'table column':
264
+ cols['table column.'+str(idx)] = (xmin, ymin-expand_rowcol_bbox_top, xmax, ymax+expand_rowcol_bbox_bottom)
265
+
266
+ idx += 1
267
+
268
+
269
+ plt.axis('on')
270
+ c2.pyplot()
271
+ return rows, cols
272
+
273
+ def sort_table_featuresv2(self, rows:dict, cols:dict):
274
+ # Sometimes the header and first row overlap, and we need the header bbox not to have first row's bbox inside the headers bbox
275
+ rows_ = {table_feature : (xmin, ymin, xmax, ymax) for table_feature, (xmin, ymin, xmax, ymax) in sorted(rows.items(), key=lambda tup: tup[1][1])}
276
+ cols_ = {table_feature : (xmin, ymin, xmax, ymax) for table_feature, (xmin, ymin, xmax, ymax) in sorted(cols.items(), key=lambda tup: tup[1][0])}
277
+
278
+ return rows_, cols_
279
+
280
+ def individual_table_featuresv2(self, pil_img, rows:dict, cols:dict):
281
+
282
+ for k, v in rows.items():
283
+ xmin, ymin, xmax, ymax = v
284
+ cropped_img = pil_img.crop((xmin, ymin, xmax, ymax))
285
+ rows[k] = xmin, ymin, xmax, ymax, cropped_img
286
+
287
+ for k, v in cols.items():
288
+ xmin, ymin, xmax, ymax = v
289
+ cropped_img = pil_img.crop((xmin, ymin, xmax, ymax))
290
+ cols[k] = xmin, ymin, xmax, ymax, cropped_img
291
+
292
+ return rows, cols
293
+
294
+
295
+ def object_to_cellsv2(self, master_row:dict, cols:dict, expand_rowcol_bbox_top, expand_rowcol_bbox_bottom, padd_left):
296
+ '''Removes redundant bbox for rows&columns and divides each row into cells from columns
297
+ Args:
298
+
299
+ Returns:
300
+
301
+
302
+ '''
303
+ cells_img = {}
304
+ header_idx = 0
305
+ row_idx = 0
306
+ previous_xmax_col = 0
307
+ new_cols = {}
308
+ new_master_row = {}
309
+ previous_ymin_row = 0
310
+ new_cols = cols
311
+ new_master_row = master_row
312
+ ## Below 2 for loops remove redundant bounding boxes ###
313
+ # for k_col, v_col in cols.items():
314
+ # xmin_col, _, xmax_col, _, col_img = v_col
315
+ # if (np.isclose(previous_xmax_col, xmax_col, atol=5)) or (xmin_col >= xmax_col):
316
+ # print('Found a column with double bbox')
317
+ # continue
318
+ # previous_xmax_col = xmax_col
319
+ # new_cols[k_col] = v_col
320
+
321
+ # for k_row, v_row in master_row.items():
322
+ # _, ymin_row, _, ymax_row, row_img = v_row
323
+ # if (np.isclose(previous_ymin_row, ymin_row, atol=5)) or (ymin_row >= ymax_row):
324
+ # print('Found a row with double bbox')
325
+ # continue
326
+ # previous_ymin_row = ymin_row
327
+ # new_master_row[k_row] = v_row
328
+ ######################################################
329
+ for k_row, v_row in new_master_row.items():
330
+
331
+ _, _, _, _, row_img = v_row
332
+ xmax, ymax = row_img.size
333
+ xa, ya, xb, yb = 0, 0, 0, ymax
334
+ row_img_list = []
335
+ # plt.imshow(row_img)
336
+ # st.pyplot()
337
+ for idx, kv in enumerate(new_cols.items()):
338
+ k_col, v_col = kv
339
+ xmin_col, _, xmax_col, _, col_img = v_col
340
+ xmin_col, xmax_col = xmin_col - padd_left - 10, xmax_col - padd_left
341
+ # plt.imshow(col_img)
342
+ # st.pyplot()
343
+ # xa + 3 : to remove borders on the left side of the cropped cell
344
+ # yb = 3: to remove row information from the above row of the cropped cell
345
+ # xb - 3: to remove borders on the right side of the cropped cell
346
+ xa = xmin_col
347
+ xb = xmax_col
348
+ if idx == 0:
349
+ xa = 0
350
+ if idx == len(new_cols)-1:
351
+ xb = xmax
352
+ xa, ya, xb, yb = xa, ya, xb, yb
353
+
354
+ row_img_cropped = row_img.crop((xa, ya, xb, yb))
355
+ row_img_list.append(row_img_cropped)
356
+
357
+ cells_img[k_row+'.'+str(row_idx)] = row_img_list
358
+ row_idx += 1
359
+
360
+ return cells_img, len(new_cols), len(new_master_row)-1
361
+
362
+ def clean_dataframe(self, df):
363
+ '''
364
+ Remove irrelevant symbols that appear with tesseractOCR
365
+ '''
366
+ # df.columns = [col.replace('|', '') for col in df.columns]
367
+
368
+ for col in df.columns:
369
+
370
+ df[col]=df[col].str.replace("'", '', regex=True)
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
+ return df
377
+
378
+ @st.cache
379
+ def convert_df(self, df):
380
+ return df.to_csv().encode('utf-8')
381
+
382
+
383
+ def create_dataframe(self, c3, cells_pytess_result:list, max_cols:int, max_rows:int):
384
+ '''Create dataframe using list of cell values of the table, also checks for valid header of dataframe
385
+ Args:
386
+ cells_pytess_result: list of strings, each element representing a cell in a table
387
+ max_cols, max_rows: number of columns and rows
388
+ Returns:
389
+ dataframe : final dataframe after all pre-processing
390
+ '''
391
+
392
+ headers = cells_pytess_result[:max_cols]
393
+ new_headers = uniquify(headers, (f' {x!s}' for x in string.ascii_lowercase))
394
+ counter = 0
395
+
396
+ cells_list = cells_pytess_result[max_cols:]
397
+ df = pd.DataFrame("", index=range(0, max_rows), columns=new_headers)
398
+
399
+ cell_idx = 0
400
+ for nrows in range(max_rows):
401
+ for ncols in range(max_cols):
402
+ df.iat[nrows, ncols] = str(cells_list[cell_idx])
403
+ cell_idx += 1
404
+
405
+ ## To check if there are duplicate headers if result of uniquify+col == col
406
+ ## This check removes headers when all headers are empty or if median of header word count is less than 6
407
+ for x, col in zip(string.ascii_lowercase, new_headers):
408
+ if f' {x!s}' == col:
409
+ counter += 1
410
+ header_char_count = [len(col) for col in new_headers]
411
+
412
+ # if (counter == len(new_headers)) or (statistics.median(header_char_count) < 6):
413
+ # st.write('woooot')
414
+ # df.columns = uniquify(df.iloc[0], (f' {x!s}' for x in string.ascii_lowercase))
415
+ # df = df.iloc[1:,:]
416
+
417
+ df = self.clean_dataframe(df)
418
+
419
+ c3.dataframe(df)
420
+ csv = self.convert_df(df)
421
+ c3.download_button("Download table", csv, "file.csv", "text/csv", key='download-csv')
422
+
423
+ return df
424
+
425
+
426
+
427
+
428
+
429
+
430
+ 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):
431
+ '''
432
+ Initiates process of generating pandas dataframes from raw pdf-page images
433
+
434
+ '''
435
+ image = Image.open(image_path).convert("RGB")
436
+ model, probas, bboxes_scaled = table_detector(image, THRESHOLD_PROBA=TD_THRESHOLD)
437
+
438
+ if bboxes_scaled.nelement() == 0:
439
+ st.write('No table found in the pdf-page image')
440
+ return ''
441
+
442
+ # try:
443
+ # st.write('Document: '+image_path.split('/')[-1])
444
+ c1, c2, c3 = st.columns((1,1,1))
445
+
446
+ self.plot_results_detection(c1, model, image, probas, bboxes_scaled, delta_xmin, delta_ymin, delta_xmax, delta_ymax)
447
+ cropped_img_list = self.crop_tables(image, probas, bboxes_scaled, delta_xmin, delta_ymin, delta_xmax, delta_ymax)
448
+
449
+ for unpadded_table in cropped_img_list:
450
+
451
+ table = self.add_padding(unpadded_table, padd_top, padd_right, padd_bottom, padd_left)
452
+ # table = super_res(table)
453
+ # table = binarizeBlur_image(table)
454
+ # table = sharpen_image(table) # Test sharpen image next
455
+ # table = td_postprocess(table)
456
+
457
+ model, probas, bboxes_scaled = table_struct_recog(table, THRESHOLD_PROBA=TSR_THRESHOLD)
458
+ rows, cols = self.generate_structure(c2, model, table, probas, bboxes_scaled, expand_rowcol_bbox_top, expand_rowcol_bbox_bottom)
459
+ # st.write(len(rows), len(cols))
460
+ rows, cols = self.sort_table_featuresv2(rows, cols)
461
+ master_row, cols = self.individual_table_featuresv2(table, rows, cols)
462
+
463
+ cells_img, max_cols, max_rows = self.object_to_cellsv2(master_row, cols, expand_rowcol_bbox_top, expand_rowcol_bbox_bottom, padd_left)
464
+
465
+ sequential_cell_img_list = []
466
+ for k, img_list in cells_img.items():
467
+ for img in img_list:
468
+ # img = super_res(img)
469
+ # img = sharpen_image(img) # Test sharpen image next
470
+ # img = binarizeBlur_image(img)
471
+ # img = self.add_padding(img, 10,10,10,10)
472
+ # plt.imshow(img)
473
+ # c3.pyplot()
474
+ sequential_cell_img_list.append(pytess(img))
475
+
476
+ cells_pytess_result = await asyncio.gather(*sequential_cell_img_list)
477
+
478
+
479
+ self.create_dataframe(c3, cells_pytess_result, max_cols, max_rows)
480
+ st.write('Errors in OCR is due to either quality of the image or performance of the OCR')
481
+ # except:
482
+ # st.write('Either incorrectly identified table or no table, to debug remove try/except')
483
+ # break
484
+ # break
485
+
486
+
487
+
488
+
489
+ if __name__ == "__main__":
490
+
491
+ img_name = st.file_uploader("Upload an image with table(s)")
492
+ st1, st2 = st.columns((1,1))
493
+ TD_th = st1.slider('Table detection threshold', 0.0, 1.0, 0.6)
494
+ TSR_th = st2.slider('Table structure recognition threshold', 0.0, 1.0, 0.8)
495
+
496
+ st1, st2, st3, st4 = st.columns((1,1,1,1))
497
+
498
+ padd_top = st1.slider('Padding top', 0, 200, 20)
499
+ padd_left = st2.slider('Padding left', 0, 200, 20)
500
+ padd_right = st3.slider('Padding right', 0, 200, 20)
501
+ padd_bottom = st4.slider('Padding bottom', 0, 200, 20)
502
+
503
+ te = TableExtractionPipeline()
504
+ # for img in image_list:
505
+ if img_name is not None:
506
+ 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))
507
+
508
+
509
+
packages.txt ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ ffmpeg
2
+ libsm6
3
+ libxext6
4
+ libgl1
5
+ tesseract-ocr-eng
6
+ python3-opencv
requirements.txt ADDED
@@ -0,0 +1,47 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ Cython==0.29.14
2
+ dask==2021.3.1
3
+ datasets==1.18.3
4
+ Flask==2.0.1
5
+ GitPython==3.1.26
6
+ imutils==0.5.4
7
+ multiprocess==0.70.12.2
8
+ numba==0.54.1
9
+ numexpr==2.7.3
10
+ numpy==1.20.3
11
+ oauthlib==3.1.0
12
+ opencv-contrib-python==4.6.0.66
13
+ openpyxl==3.0.7
14
+ Pillow==9.0.1
15
+ plotly==4.14.3
16
+ ply==3.11
17
+ protobuf==3.14.0
18
+ psutil==5.8.0
19
+ pyarrow==7.0.0
20
+ pydantic==1.7.3
21
+ pydeck==0.7.1
22
+ PyDictionary==2.0.1
23
+ pydot==1.4.2
24
+ pymongo==4.0.2
25
+ Pympler==1.0.1
26
+ PyMuPDF==1.20.2
27
+ pyperclip==1.8.2
28
+ pyppeteer==0.2.5
29
+ pyquery==1.4.3
30
+ pyreadline3==3.3
31
+ pytesseract==0.3.10
32
+ pytz-deprecation-shim==0.1.0.post0
33
+ PyWavelets==1.1.1
34
+ PyYAML==5.4.1
35
+ scipy==1.4.1
36
+ seaborn==0.11.1
37
+ sklearn==0.0
38
+ streamlit==1.5.1
39
+ timm==0.6.7
40
+ tokenizers==0.12.1
41
+ toml==0.10.2
42
+ toolz==0.11.1
43
+ torch==1.10.0
44
+ torchvision==0.11.1
45
+ git+https://github.com/huggingface/transformers.git
46
+ #-e git+https://github.com/nielsrogge/transformers.git@d34f7e6ffbb911d39465173ef2b35ba147ef58a9#egg=transformers
47
+ urllib3==1.26.7