change vietocr to PaddleOCR

README.md

@@ -1,8 +1,8 @@
 ---
-title: Table Extraction
+title: Table Extraction (Table Transformer + PaddleOCR)
 emoji: 🚀
-colorFrom: indigo
-colorTo: purple
+colorFrom: while
+colorTo: black
 sdk: streamlit
 sdk_version: 1.21.0
 app_file: app.py
@@ -12,7 +12,6 @@ pinned: false
 
 Imported from https://huggingface.co/spaces/jurgendn/table-extraction with some adjustment.
 
-
 Current pipeline:
 
 Table detection: https://huggingface.co/microsoft/table-transformer-detection
@@ -20,3 +19,5 @@ Table detection: https://huggingface.co/microsoft/table-transformer-detection
 Table recognition: https://huggingface.co/microsoft/table-transformer-structure-recognition
 
 OCR: https://github.com/pbcquoc/vietocr
+
+OCR-new: https://github.com/PaddlePaddle/PaddleOCR

app.py

@@ -1,5 +1,6 @@
 import asyncio
 import string
+import random
 from collections import Counter
 from itertools import count, tee
 
@@ -11,8 +12,9 @@ import streamlit as st
 import torch
 from PIL import Image
 from transformers import DetrImageProcessor, TableTransformerForObjectDetection
-from vietocr.tool.config import Cfg
-from vietocr.tool.predictor import Predictor
+from paddleocr import PaddleOCR
+
+ocr = PaddleOCR(use_angle_cls=True, lang="en",use_gpu=False)
 
 st.set_option('deprecation.showPyplotGlobalUse', False)
 st.set_page_config(layout='wide')
@@ -20,13 +22,6 @@ st.title("Table Detection and Table Structure Recognition")
 st.write(
     "Implemented by MSFT team: https://github.com/microsoft/table-transformer")
 
-# config = Cfg.load_config_from_name('vgg_transformer')
-config = Cfg.load_config_from_name('vgg_seq2seq')
-config['cnn']['pretrained'] = False
-config['device'] = 'cpu'
-config['predictor']['beamsearch'] = False
-detector = Predictor(config)
-
 table_detection_model = TableTransformerForObjectDetection.from_pretrained(
     "microsoft/table-transformer-detection")
 
@@ -43,10 +38,20 @@ def cv_to_PIL(cv_img):
 
 
 async def pytess(cell_pil_img, threshold: float = 0.5):
-    text, prob = detector.predict(cell_pil_img, return_prob=True)
-    if prob < threshold:
-        return ""
-    return text.strip()
+    cell_pil_img = TableExtractionPipeline.add_padding(pil_img=cell_pil_img, top=50, right=30, bottom=50, left=30, color=(255, 255, 255))
+    result = ocr.ocr(np.asarray(cell_pil_img), cls=True)[0]
+    
+    #Debug
+    # filename = str(random.random())
+    # cell_pil_img.save("dump/" + filename + ".png")
+    # print(filename)
+    # print(result)
+    
+    text = ""
+    if result != None:
+        txts = [line[1][0] for line in result]
+        text = " ".join(txts)
+    return text
 
 
 def sharpen_image(pil_img):
@@ -179,8 +184,8 @@ class TableExtractionPipeline():
 
     # colors = ["red", "blue", "green", "red", "red", "red"]
 
-    def add_padding(self,
-                    pil_img,
+    @staticmethod
+    def add_padding(pil_img,
                     top,
                     right,
                     bottom,
@@ -449,11 +454,17 @@ class TableExtractionPipeline():
 
         c3.dataframe(df)
         csv = self.convert_df(df)
+
+        try:
+            numkey = df.iloc[0, 0]
+        except:
+            numkey = str(0)
+
         c3.download_button("Download table",
                            csv,
                            "file.csv",
                            "text/csv",
-                           key='download-csv-' + df.iloc[0, 0])
+                           key='download-csv-' + numkey)
 
         return df
 

requirements.txt

@@ -6,4 +6,6 @@ vietocr==0.3.11
 streamlit==1.21.0
 pandas
 transformers==4.29.1
-Pillow==9.5.0
+Pillow==9.5.0
+paddlepaddle
+paddleocr