Spaces:
Build error
Build error
Duplicate from king007/TableTransformer2CSV
Browse filesCo-authored-by: king007 <king007@users.noreply.huggingface.co>
- .gitattributes +31 -0
- README.md +13 -0
- app.py +510 -0
- packages.txt +6 -0
- 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: king007/TableTransformer2CSV
|
11 |
+
---
|
12 |
+
|
13 |
+
Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
|
app.py
ADDED
@@ -0,0 +1,510 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
+
|
33 |
+
def PIL_to_cv(pil_img):
|
34 |
+
return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
|
35 |
+
|
36 |
+
def cv_to_PIL(cv_img):
|
37 |
+
return Image.fromarray(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB))
|
38 |
+
|
39 |
+
|
40 |
+
async 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 docs
|
47 |
+
# '''
|
48 |
+
# requires opencv-contrib-python installed without the opencv-python
|
49 |
+
# 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 = 2
|
54 |
+
# 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_img
|
60 |
+
|
61 |
+
|
62 |
+
def 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_img
|
72 |
+
|
73 |
+
|
74 |
+
def 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-list
|
77 |
+
`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 |
+
continue
|
88 |
+
else:
|
89 |
+
seq[idx] += suffix
|
90 |
+
|
91 |
+
return seq
|
92 |
+
|
93 |
+
def 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 - result
|
99 |
+
return cv_to_PIL(result)
|
100 |
+
|
101 |
+
|
102 |
+
|
103 |
+
def td_postprocess(pil_img):
|
104 |
+
'''
|
105 |
+
Removes gray background from tables
|
106 |
+
'''
|
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 & nzmask
|
114 |
+
|
115 |
+
new_img = img.copy()
|
116 |
+
new_img[np.where(mask)] = 255
|
117 |
+
|
118 |
+
|
119 |
+
return cv_to_PIL(new_img)
|
120 |
+
|
121 |
+
# def super_res(pil_img):
|
122 |
+
# # requires opencv-contrib-python installed without the opencv-python
|
123 |
+
# 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_img
|
135 |
+
|
136 |
+
def table_detector(image, THRESHOLD_PROBA):
|
137 |
+
'''
|
138 |
+
Table detection using DEtect-object TRansformer pre-trained on 1 million tables
|
139 |
+
|
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_PROBA
|
152 |
+
|
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 |
+
|
160 |
+
def table_struct_recog(image, THRESHOLD_PROBA):
|
161 |
+
'''
|
162 |
+
Table structure recognition using DEtect-object TRansformer pre-trained on 1 million tables
|
163 |
+
'''
|
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_PROBA
|
174 |
+
|
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 |
+
|
185 |
+
class 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 edges
|
194 |
+
'''
|
195 |
+
width, height = pil_img.size
|
196 |
+
new_width = width + right + left
|
197 |
+
new_height = height + top + bottom
|
198 |
+
result = Image.new(pil_img.mode, (new_width, new_height), color)
|
199 |
+
result.paste(pil_img, (left, top))
|
200 |
+
return result
|
201 |
+
|
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_list
|
235 |
+
|
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 table
|
240 |
+
'''
|
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 = 0
|
249 |
+
|
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 += 1
|
268 |
+
|
269 |
+
|
270 |
+
plt.axis('on')
|
271 |
+
c2.pyplot()
|
272 |
+
return rows, cols
|
273 |
+
|
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 bbox
|
276 |
+
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 = v
|
285 |
+
cropped_img = pil_img.crop((xmin, ymin, xmax, ymax))
|
286 |
+
rows[k] = xmin, ymin, xmax, ymax, cropped_img
|
287 |
+
|
288 |
+
for k, v in cols.items():
|
289 |
+
xmin, ymin, xmax, ymax = v
|
290 |
+
cropped_img = pil_img.crop((xmin, ymin, xmax, ymax))
|
291 |
+
cols[k] = xmin, ymin, xmax, ymax, cropped_img
|
292 |
+
|
293 |
+
return rows, cols
|
294 |
+
|
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 columns
|
298 |
+
Args:
|
299 |
+
|
300 |
+
Returns:
|
301 |
+
|
302 |
+
|
303 |
+
'''
|
304 |
+
cells_img = {}
|
305 |
+
header_idx = 0
|
306 |
+
row_idx = 0
|
307 |
+
previous_xmax_col = 0
|
308 |
+
new_cols = {}
|
309 |
+
new_master_row = {}
|
310 |
+
previous_ymin_row = 0
|
311 |
+
new_cols = cols
|
312 |
+
new_master_row = master_row
|
313 |
+
## 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_col
|
316 |
+
# if (np.isclose(previous_xmax_col, xmax_col, atol=5)) or (xmin_col >= xmax_col):
|
317 |
+
# print('Found a column with double bbox')
|
318 |
+
# continue
|
319 |
+
# previous_xmax_col = xmax_col
|
320 |
+
# new_cols[k_col] = v_col
|
321 |
+
|
322 |
+
# for k_row, v_row in master_row.items():
|
323 |
+
# _, ymin_row, _, ymax_row, row_img = v_row
|
324 |
+
# if (np.isclose(previous_ymin_row, ymin_row, atol=5)) or (ymin_row >= ymax_row):
|
325 |
+
# print('Found a row with double bbox')
|
326 |
+
# continue
|
327 |
+
# previous_ymin_row = ymin_row
|
328 |
+
# new_master_row[k_row] = v_row
|
329 |
+
######################################################
|
330 |
+
for k_row, v_row in new_master_row.items():
|
331 |
+
|
332 |
+
_, _, _, _, row_img = v_row
|
333 |
+
xmax, ymax = row_img.size
|
334 |
+
xa, ya, xb, yb = 0, 0, 0, ymax
|
335 |
+
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 = kv
|
340 |
+
xmin_col, _, xmax_col, _, col_img = v_col
|
341 |
+
xmin_col, xmax_col = xmin_col - padd_left - 10, xmax_col - padd_left
|
342 |
+
# plt.imshow(col_img)
|
343 |
+
# st.pyplot()
|
344 |
+
# xa + 3 : to remove borders on the left side of the cropped cell
|
345 |
+
# yb = 3: to remove row information from the above row of the cropped cell
|
346 |
+
# xb - 3: to remove borders on the right side of the cropped cell
|
347 |
+
xa = xmin_col
|
348 |
+
xb = xmax_col
|
349 |
+
if idx == 0:
|
350 |
+
xa = 0
|
351 |
+
if idx == len(new_cols)-1:
|
352 |
+
xb = xmax
|
353 |
+
xa, ya, xb, yb = xa, ya, xb, yb
|
354 |
+
|
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_list
|
359 |
+
row_idx += 1
|
360 |
+
|
361 |
+
return cells_img, len(new_cols), len(new_master_row)-1
|
362 |
+
|
363 |
+
def clean_dataframe(self, df):
|
364 |
+
'''
|
365 |
+
Remove irrelevant symbols that appear with tesseractOCR
|
366 |
+
'''
|
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 df
|
378 |
+
|
379 |
+
@st.cache
|
380 |
+
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 dataframe
|
386 |
+
Args:
|
387 |
+
cells_pytess_result: list of strings, each element representing a cell in a table
|
388 |
+
max_cols, max_rows: number of columns and rows
|
389 |
+
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 = 0
|
396 |
+
|
397 |
+
cells_list = cells_pytess_result[max_cols:]
|
398 |
+
df = pd.DataFrame("", index=range(0, max_rows), columns=new_headers)
|
399 |
+
|
400 |
+
cell_idx = 0
|
401 |
+
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 += 1
|
405 |
+
|
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 6
|
408 |
+
for x, col in zip(string.ascii_lowercase, new_headers):
|
409 |
+
if f' {x!s}' == col:
|
410 |
+
counter += 1
|
411 |
+
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 df
|
425 |
+
|
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 images
|
434 |
+
|
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 next
|
456 |
+
# 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 next
|
471 |
+
# 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 |
+
# break
|
485 |
+
# break
|
486 |
+
|
487 |
+
|
488 |
+
|
489 |
+
|
490 |
+
if __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 |
+
|
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
|