nielsr HF staff commited on
Commit
f29cd95
1 Parent(s): 8dd7ef6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +13 -2
app.py CHANGED
@@ -4,6 +4,7 @@ from matplotlib.patches import Patch
4
  import io
5
  from PIL import Image, ImageDraw
6
  import numpy as np
 
7
 
8
  from transformers import TableTransformerImageProcessor, AutoModelForObjectDetection
9
  import torch
@@ -220,7 +221,17 @@ def apply_ocr(cell_coordinates, cropped_table):
220
  row_data = row_data + ["" for _ in range(max_num_columns - len(row_data))]
221
  data[row] = row_data
222
 
223
- return data
 
 
 
 
 
 
 
 
 
 
224
 
225
 
226
  def process_pdf(image):
@@ -241,7 +252,7 @@ examples =[['image.png']]
241
 
242
  app = gr.Interface(fn=process_pdf,
243
  inputs=gr.Image(type="pil"),
244
- outputs=[gr.Image(type="pil", label="Detected table"), "json"],
245
  title=title,
246
  description=description,
247
  examples=examples)
 
4
  import io
5
  from PIL import Image, ImageDraw
6
  import numpy as np
7
+ import csv
8
 
9
  from transformers import TableTransformerImageProcessor, AutoModelForObjectDetection
10
  import torch
 
221
  row_data = row_data + ["" for _ in range(max_num_columns - len(row_data))]
222
  data[row] = row_data
223
 
224
+ # write to csv
225
+ with open('output.csv','w') as result_file:
226
+ wr = csv.writer(result_file, dialect='excel')
227
+
228
+ for row, row_text in data.items():
229
+ wr.writerow(row_text)
230
+
231
+ # return as Pandas dataframe
232
+ df = pd.read_csv('output.csv')
233
+
234
+ return df
235
 
236
 
237
  def process_pdf(image):
 
252
 
253
  app = gr.Interface(fn=process_pdf,
254
  inputs=gr.Image(type="pil"),
255
+ outputs=[gr.Image(type="pil", label="Detected table"), "dataframe"],
256
  title=title,
257
  description=description,
258
  examples=examples)