upgrade to GPU; move to taskflow

app.py

@@ -1,5 +1,5 @@
 #-*- coding: UTF-8 -*-
-# Copyright 2022 The Impira Team and the HuggingFace Team.
+# Copyright 2022 the HuggingFace Team.
 # Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,21 +14,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-import json
 import base64
-from io import BytesIO
-from PIL import Image
 import traceback
 
-import requests
-import numpy as np
 import gradio as gr
 import cv2
 
+from paddlenlp import Taskflow
 from paddlenlp.utils.doc_parser import DocParser
 
 doc_parser = DocParser()
+task_instance = Taskflow("information_extraction", model="uie-x-base", schema="")
 
 examples = [
     [
@@ -38,7 +34,7 @@ examples = [
     [
         "license.jpeg",
         "Name;DOB;ISS;EXP",
-    ],
+],
     [
         "invoice.jpeg",
         "名称;纳税人识别号;开票日期",
@@ -62,11 +58,11 @@ example_files = {
 }
 
 lang_map = {
-    "resume.png": "ch-no",
-    "custom.jpeg": "ch-no",
-    "business_card.png": "en-no",
-    "invoice.jpeg": "ch-no",
-    "license.jpeg": "en-no",
+    "resume.png": "ch",
+    "custom.jpeg": "ch",
+    "business_card.png": "en",
+    "invoice.jpeg": "ch",
+    "license.jpeg": "en",
 }
 
 def dbc2sbc(s):
@@ -139,43 +135,73 @@ def np2base64(image_np):
     return base64_str
 
 
+def get_schema(schema_str):
+    def _is_ch(s):
+        for ch in s:
+            if "\u4e00" <= ch <= "\u9fff":
+                return True
+        return False
+    schema_lang = "ch" if _is_ch(schema_str) else "en"
+    schema = schema_str.split(";")
+    schema_list = []
+    for s in schema:
+        cand = s.split("|")
+        if len(cand) == 1:
+            schema_list.append(cand[0])
+        else:
+            subject = cand[0]
+            relations = cand[1:]
+            added = False
+            for a in schema_list:
+                if isinstance(a, dict):
+                    if subject in a.keys():
+                        a[subject].extend(relations)
+                        added = True
+                        break
+            if not added:
+                a = {subject: relations}
+                schema_list.append(a)
+    return schema_list, schema_lang
+
+
+def run_taskflow(document, schema, argument):
+    task_instance.set_schema(schema)
+    # task_instance.set_argument(argument)
+    return task_instance({'doc': document})
+
+
 def process_doc(document, schema, ocr_lang, layout_analysis):
     if not schema:
         schema = '时间;组织机构;人物'
     if document is None:
         return None, None
 
-    option = ocr_lang + "-" + layout_analysis
-
-    schema = dbc2sbc(schema)
-
-    access_token = os.environ['token']
-    url = f"https://aip.baidubce.com/rpc/2.0/nlp-itec/poc/ie?access_token={access_token}"
-
-    base64_str = np2base64(doc_parser.read_image(document))
-
-    r = requests.post(url, json={"doc": base64_str, "schema": schema, "option": option})
-    response = r.json()
-    print(response)    
-    predictions = response['result']
+    schema, schema_lang = get_schema(dbc2sbc(schema))
+    argument = {
+        "ocr_lang": ocr_lang,
+        "schema_lang": schema_lang,
+        "layout_analysis": layout_analysis
+    }
+    prediction = run_taskflow(document, schema, argument)[0]
 
     img_show = doc_parser.write_image_with_results(
-        base64_str,
-        result=predictions, 
-        max_size=2000, 
+        document,
+        result=prediction,
         return_image=True)
     img_list = [img_show]
 
     return (
         gr.update(visible=True, value=img_list),
-        gr.update(visible=True, value=predictions),
+        gr.update(visible=True, value=prediction),
     )
 
 
 def load_example_document(img, schema, ocr_lang, layout_analysis):
     if img is not None:
         document = example_files[schema]
-        ocr_lang, layout_analysis = lang_map[document].split("-")
+        choice = lang_map[document].split("-")
+        ocr_lang = choice[0]
+        layout_analysis = False if len(choice) == 1 else True
         preview, answer = process_doc(document, schema, ocr_lang, layout_analysis)
         return document, schema, preview, gr.update(visible=True), answer
     else:

requirements.txt

@@ -1,5 +1,6 @@
-numpy
+numpy==1.21.6
 opencv-python
-paddlenlp
-paddlepaddle
-requests
+# install from git for now, wait until v2.4.6 release
+PaddleNLP
+paddleocr
+paddlepaddle-gpu