[App] Filter low confidence ceil

app.py

@@ -25,7 +25,7 @@ st.write(
 config = Cfg.load_config_from_name('vgg_seq2seq')
 config['cnn']['pretrained'] = False
 config['device'] = 'cpu'
-config['predictor']['beamsearch'] = True
+config['predictor']['beamsearch'] = False
 detector = Predictor(config)
 
 table_detection_model = TableTransformerForObjectDetection.from_pretrained(
@@ -43,8 +43,11 @@ def cv_to_PIL(cv_img):
     return Image.fromarray(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB))
 
 
-async def pytess(cell_pil_img):
-    text = detector.predict(cell_pil_img)
+async def pytess(cell_pil_img, threshold: float = 0.5):
+    text, prob = detector.predict(cell_pil_img, return_prob=True)
+    st.write(prob)
+    if prob < threshold:
+        return ""
     return text.strip()
 
 
@@ -457,9 +460,10 @@ class TableExtractionPipeline():
         return df
 
     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):
+                            OCR_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):
         '''
         Initiates process of generating pandas dataframes from raw pdf-page images
 
@@ -516,7 +520,8 @@ class TableExtractionPipeline():
                     # img = self.add_padding(img, 10,10,10,10)
                     # plt.imshow(img)
                     # c3.pyplot()
-                    sequential_cell_img_list.append(pytess(img))
+                    sequential_cell_img_list.append(
+                        pytess(cell_pil_img=img, threshold=OCR_THRESHOLD))
 
             cell_ocr_res = await asyncio.gather(*sequential_cell_img_list)
 
@@ -533,9 +538,10 @@ class TableExtractionPipeline():
 if __name__ == "__main__":
 
     img_name = st.file_uploader("Upload an image with table(s)")
-    st1, st2 = st.columns((1, 1))
+    st1, st2, st3 = st.columns((1, 1, 1))
     TD_th = st1.slider('Table detection threshold', 0.0, 1.0, 0.8)
     TSR_th = st2.slider('Table structure recognition threshold', 0.0, 1.0, 0.8)
+    OCR_th = st3.slider("Text Probs Threshold", 0.0, 1.0, 0.5)
 
     st1, st2, st3, st4 = st.columns((1, 1, 1, 1))
 
@@ -551,6 +557,7 @@ if __name__ == "__main__":
             te.start_process(img_name,
                              TD_THRESHOLD=TD_th,
                              TSR_THRESHOLD=TSR_th,
+                             OCR_THRESHOLD=OCR_th,
                              padd_top=padd_top,
                              padd_left=padd_left,
                              padd_bottom=padd_bottom,