Update app.py

app.py

@@ -4,12 +4,15 @@ import gradio as gr
 # from concurrent.futures import ThreadPoolExecutor
 import pdfplumber
 import pandas as pd
+import time
+from cnocr import CnOcr
 from sentence_transformers import SentenceTransformer, models, util
 word_embedding_model = models.Transformer('uer/sbert-base-chinese-nli', do_lower_case=True)
 pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(), pooling_mode='cls')
 embedder = SentenceTransformer(modules=[word_embedding_model, pooling_model])
-
-url = 'https://souljoy-my-api.hf.space/qa_maker'
+ocr = CnOcr()
+# chat_url = 'https://souljoy-my-api.hf.space/sale'
+chat_url = 'https://souljoy-my-api.hf.space/chatpdf'
 headers = {
     'Content-Type': 'application/json',
 }
@@ -41,7 +44,7 @@ def doc_emb(doc: str):
     #     emb_list.append(f.result())
     print('\n'.join(texts))
     return texts, emb_list, gr.Textbox.update(visible=True), gr.Button.update(visible=True), gr.Markdown.update(
-        value="""操作说明 step 3：PDF解析提交成功 🙋 可以开始对话啦~"""), gr.Chatbot.update(visible=True)
+        value="""操作说明 step 3：PDF解析提交成功！ 🙋 可以开始对话啦~"""), gr.Chatbot.update(visible=True)
 
 
 def get_response(msg, bot, doc_text_list, doc_embeddings):
@@ -89,7 +92,7 @@ def get_response(msg, bot, doc_text_list, doc_embeddings):
     req_json['doc'] = '' if len(sub_doc_list) == 0 else '\n'.join(sub_doc_list)
     data = {"content": json.dumps(req_json)}
     print('data:\n', req_json)
-    result = requests.post(url='https://souljoy-my-api.hf.space/chatpdf',
+    result = requests.post(url=chat_url,
                            data=json.dumps(data),
                            headers=headers
                            )
@@ -107,6 +110,17 @@ def up_file(files):
                 # 读取PDF文档第i+1页
                 page = pdf.pages[i]
                 res_list = page.extract_text().split('\n')[:-1]
+
+                for j in range(len(page.images)):
+                    # 获取图片的二进制流
+                    img = page.images[j]
+                    file_name = '{}-{}-{}.png'.format(str(time.time()), str(i), str(j))
+                    with open(file_name, mode='wb') as f:
+                        f.write(img['stream'].get_data())
+                    res = ocr.ocr(file_name)
+                    if len(res) > 0:
+                        res_list.append(' '.join([re['text'] for re in res]))
+
                 tables = page.extract_tables()
                 for table in tables:
                     # 第一列当成表头：
@@ -124,7 +138,7 @@ def up_file(files):
         print(i)
     return gr.Textbox.update(value='\n'.join(doc_text_list), visible=True), gr.Button.update(
         visible=True), gr.Markdown.update(
-        value="操作说明 step 2：确认PDF解析结果（可修正），点击“提交结果”，进行对话")
+        value="操作说明 step 2：确认PDF解析结果（可修正），点击“提交解析结果”，随后进行对话")
 
 
 with gr.Blocks() as demo: