initial space

apis/layoutlm.py

@@ -0,0 +1,84 @@
+from transformers import pipeline
+import pandas as pd
+
+class LayoutLM:
+
+    def __init__(self, save_pretrained_fpath:str=None) -> None:
+        self.pipeline_category = 'document-question-answering'
+        self.tf_pipeline = pipeline
+        self.pipeline = None
+
+        if save_pretrained_fpath is not None:
+            pipe = self.tf_pipeline(self.pipeline_category)
+            pipe.save_pretrained(save_pretrained_fpath)
+
+        self.default_model = 'impira/layoutlm-invoices'
+        self.default_ex_answer = {'score':0, 'answer':'-'}
+
+    def set_model(self, model:str):
+        if model is None:
+            model = self.default_model
+
+        self.pipeline = self.tf_pipeline(self.pipeline_category, model=model)
+
+    def answer_the_question(self, img, question: str, is_debug=False):
+        score = 0
+        answer = '-'
+        answers = None
+        if self.pipeline is not None:
+            answers = self.pipeline(img, question)
+
+            for a in answers:
+                if a['score'] > score:
+                    score = a['score']
+                    answer = a['answer']
+
+        if is_debug:
+            print('--------------------')
+            print(f'Q: {question}\nA: {answer} (acc:{score:.2f})\n')
+            print(answers)
+
+        return answer
+
+    def inference(self, img, is_debug=False):
+        merchant_id = self.answer_the_question(img, 'What is merchant ID?')
+        merchant_name = self.answer_the_question(img, 'What is merchant name?')
+        merchant_address = self.answer_the_question(img, 'What is merchant address?')
+        invoice_no = self.answer_the_question(img, 'What is tax invoice number?')
+        products = self.answer_the_question(img, 'What are buy products?')
+        pos_no = self.answer_the_question(img, 'What is POS number?')
+        net_price = self.answer_the_question(img, 'What is the net-price?')
+        date_time = self.answer_the_question(img, 'What date, year and time of the invoice?')
+
+        if is_debug:
+            print(f'Merchant ID: {merchant_id}')
+            print(f'Merchant name: {merchant_name}')
+            print(f'Merchant address: {merchant_address}')
+            print(f'Invoice no.: {invoice_no}')
+            print(f'Products: {products}')
+            print(f'POS no.: {pos_no}')
+            print(f'Net price: {net_price}')
+            print(f'Date/Time: {date_time}')
+
+        return pd.DataFrame({
+            'Data' : [
+                'Merchant ID', 
+                'Merchant name', 
+                'Merchant address', 
+                'Invoice no.',
+                'Products',
+                'POS no.',
+                'Net price',
+                'Date/Time'
+            ], 
+            'Value' : [
+                str(merchant_id),
+                str(merchant_name),
+                str(merchant_address),
+                str(invoice_no),
+                str(products),
+                str(pos_no),
+                str(net_price),
+                str(date_time)
+            ]
+        })

demo.py

@@ -0,0 +1,36 @@
+from apis.layoutlm import LayoutLM
+import pandas as pd
+import gradio as gr
+import os
+
+layoutlm = None
+
+def auth(username, password):
+    u = os.environ.get('USERNAME')
+    p = os.environ.get('PASSWORD')
+    return (username == u and password == p)
+
+def inference(img) -> pd.DataFrame:
+    return layoutlm.inference(img)
+
+if __name__ == "__main__":
+    try:
+        layoutlm = LayoutLM()
+        layoutlm.set_model(layoutlm.default_model)
+
+        demo = gr.Interface(
+            inference,
+            gr.Image(),
+            gr.Dataframe(
+                headers=['Data', 'Value'],
+                datatype=['str', 'str'],
+                row_count=8,
+                col_count=(2, 'fixed'), 
+                interactive=False
+            )
+        )
+
+        demo.launch(auth=auth)
+
+    except Exception as e:
+        print(str(e))