Update app.py

app.py

@@ -16,7 +16,8 @@ import numpy as np
 # from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 # from cv2 import dnn_superres
 from transformers import DetrFeatureExtractor
-from transformers import DetrForObjectDetection
+#from transformers import DetrForObjectDetection
+from transformers import TableTransformerForObjectDetection
 import torch
 import asyncio
 # pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'
@@ -36,11 +37,8 @@ def cv_to_PIL(cv_img):
 
 
 async def pytess(cell_pil_img):
-    pytess_output=' '.join(pytesseract.image_to_data(cell_pil_img, output_type=Output.DICT, config='-c tessedit_char_blacklist=œ˜â€œï¬â™Ã©œ¢!|”?«“¥ --psm 6 preserve_interword_spaces')['text']).strip()
-    print("pytess_output#####################################################")
-    print(pytess_output)
-    print("pytess_output@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@")
-    return pytess_output
+    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()
+
 
 # def super_res(pil_img):
     # '''
@@ -142,7 +140,7 @@ def table_detector(image, THRESHOLD_PROBA):
     feature_extractor = DetrFeatureExtractor(do_resize=True, size=800, max_size=800)
     encoding = feature_extractor(image, return_tensors="pt")
 
-    model = DetrForObjectDetection.from_pretrained("microsoft/table-transformer-detection")
+    model = TableTransformerForObjectDetection.from_pretrained("microsoft/table-transformer-detection")
 
     with torch.no_grad():
         outputs = model(**encoding)
@@ -165,7 +163,7 @@ def table_struct_recog(image, THRESHOLD_PROBA):
     feature_extractor = DetrFeatureExtractor(do_resize=True, size=1000, max_size=1000)
     encoding = feature_extractor(image, return_tensors="pt")
 
-    model = DetrForObjectDetection.from_pretrained("microsoft/table-transformer-structure-recognition")
+    model = TableTransformerForObjectDetection.from_pretrained("microsoft/table-transformer-structure-recognition")
     with torch.no_grad():
         outputs = model(**encoding)
 
@@ -432,11 +430,9 @@ class TableExtractionPipeline():
         '''
         image = Image.open(image_path).convert("RGB")
         model, probas, bboxes_scaled = table_detector(image, THRESHOLD_PROBA=TD_THRESHOLD)
-        st.text(bboxes_scaled)
-        st.image(image, caption='input')
-        st.text(bboxes_scaled.nelement())
+
         if bboxes_scaled.nelement() == 0:
-            print('No table found in the pdf-page image'+image_path.split('/')[-1])
+            st.write('No table found in the pdf-page image')
             return ''
         
         # try:
@@ -489,17 +485,21 @@ class TableExtractionPipeline():
 if __name__ == "__main__":
 
     img_name = st.file_uploader("Upload an image with table(s)")
+    st1, st2 = st.columns((1,1))
+    TD_th = st1.slider('Table detection threshold', 0.0, 1.0, 0.6)
+    TSR_th = st2.slider('Table structure recognition threshold', 0.0, 1.0, 0.8)
 
-    padd_top = st.slider('Padding top', 0, 200, 20)
-    padd_left = st.slider('Padding left', 0, 200, 20)
-    padd_right = st.slider('Padding right', 0, 200, 20)
-    padd_bottom = st.slider('Padding bottom', 0, 200, 20)
+    st1, st2, st3, st4 = st.columns((1,1,1,1))
 
+    padd_top = st1.slider('Padding top', 0, 200, 20)
+    padd_left = st2.slider('Padding left', 0, 200, 20)
+    padd_right = st3.slider('Padding right', 0, 200, 20)
+    padd_bottom = st4.slider('Padding bottom', 0, 200, 20)
 
     te = TableExtractionPipeline()
     # for img in image_list:
     if img_name is not None:
-        asyncio.run(te.start_process(img_name, TD_THRESHOLD=0.6, TSR_THRESHOLD=0.8, 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))
+        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))