Jiang Xiaolan commited on
Commit
23a6a28
·
1 Parent(s): fbaa6c2
Files changed (2) hide show
  1. app.py +140 -0
  2. requirements.txt +3 -0
app.py ADDED
@@ -0,0 +1,140 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import sys
3
+ import subprocess
4
+ import streamlit as st
5
+ import io
6
+ import pypdfium2
7
+ from PIL import Image
8
+ import logging
9
+
10
+ # 设置日志记录器
11
+ logging.basicConfig(level=logging.INFO)
12
+ logger = logging.getLogger(__name__)
13
+
14
+
15
+ def clone_repo():
16
+ # 从环境变量中获取 GitHub Token
17
+ github_token = os.getenv('GH_TOKEN')
18
+
19
+ if github_token is None:
20
+ logger.error("GitHub token is not set. Please set the GH_TOKEN secret in your Space settings.")
21
+ return False
22
+
23
+ # 使用 GitHub Token 进行身份验证并克隆仓库
24
+ clone_command = f'git clone https://{github_token}@github.com/mamba-ai/invoice_agent.git'
25
+ if os.path.exists('invoice_agent'):
26
+ logger.warning("Repository already exists.")
27
+ repo_dir = 'invoice_agent'
28
+ # 将仓库路径添加到 Python 模块搜索路径中
29
+ logger.warning(f"Adding {os.path.abspath(repo_dir)} to sys.path")
30
+ sys.path.append(os.path.abspath(repo_dir))
31
+ return True
32
+ else:
33
+ logger.info("Cloning repository...")
34
+ result = subprocess.run(clone_command, shell=True, capture_output=True, text=True)
35
+
36
+ if result.returncode == 0:
37
+ logger.warning("Repository cloned successfully.")
38
+ repo_dir = 'invoice_agent'
39
+
40
+ # 将仓库路径添加到 Python 模块搜索路径中
41
+ sys.path.append(os.path.abspath(repo_dir))
42
+ logger.warning(f"Adding {os.path.abspath(repo_dir)} to sys.path")
43
+ return True
44
+ else:
45
+ logger.error(f"Failed to clone repository: {result.stderr}")
46
+ return False
47
+
48
+
49
+
50
+ if clone_repo():
51
+ # 克隆成功后导入模块
52
+ import invoice_agent.agent as ia
53
+ # from invoice_agent.agent import load_models, get_ocr_predictions, get_json_result
54
+
55
+ def open_pdf(pdf_file):
56
+ stream = io.BytesIO(pdf_file.getvalue())
57
+ return pypdfium2.PdfDocument(stream)
58
+
59
+
60
+ @st.cache_data()
61
+ def get_page_image(pdf_file, page_num, dpi=96):
62
+ doc = open_pdf(pdf_file)
63
+ renderer = doc.render(
64
+ pypdfium2.PdfBitmap.to_pil,
65
+ page_indices=[page_num - 1],
66
+ scale=dpi / 72,
67
+ )
68
+ png = list(renderer)[0]
69
+ png_image = png.convert("RGB")
70
+ return png_image
71
+
72
+
73
+ @st.cache_data()
74
+ def page_count(pdf_file):
75
+ doc = open_pdf(pdf_file)
76
+ return len(doc)
77
+
78
+ st.set_page_config(layout="wide")
79
+
80
+ models = ia.load_models()
81
+
82
+ st.title("""
83
+ 受領した請求書を自動で電子化 (Demo)
84
+ """)
85
+
86
+ col1, _, col2 = st.columns([.45, 0.1, .45])
87
+
88
+ in_file = st.sidebar.file_uploader(
89
+ "PDFファイルまたは画像:",
90
+ type=["pdf", "png", "jpg", "jpeg", "gif", "webp"],
91
+ )
92
+
93
+ if in_file is None:
94
+ st.stop()
95
+
96
+ filetype = in_file.type
97
+ whole_image = False
98
+ if "pdf" in filetype:
99
+ page_count = page_count(in_file)
100
+ page_number = st.sidebar.number_input(f"ページ番号 {page_count}:", min_value=1, value=1, max_value=page_count)
101
+
102
+ pil_image = get_page_image(in_file, page_number)
103
+ else:
104
+ pil_image = Image.open(in_file).convert("RGB")
105
+
106
+ text_rec = st.sidebar.button("認識開始")
107
+
108
+ if pil_image is None:
109
+ st.stop()
110
+
111
+ with col1:
112
+ st.write("## アップロードされたファイル")
113
+ st.image(pil_image, caption="アップロードされたファイル", use_column_width=True)
114
+
115
+ if text_rec:
116
+ with col2:
117
+ st.write("## 結果")
118
+
119
+ # Placeholder for status indicator
120
+ status_placeholder = st.empty()
121
+
122
+ with st.spinner('現在ファイルを解析中です'):
123
+ # Simulate model running time
124
+ # time.sleep(5) # Replace this with actual model running code
125
+ predictions = ia.get_ocr_predictions(pil_image, models)
126
+
127
+ # Simulate OCR result as a JSON object
128
+ json_predictions = ia.get_json_result(predictions)
129
+
130
+ # After model finishes
131
+ status_placeholder.success('ファイルの解析が完了しました!')
132
+
133
+ # Display the result
134
+ st.write("解析後の内容:")
135
+ st.json(json_predictions)
136
+ # st.write(predictions)
137
+
138
+
139
+
140
+
requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ surya-ocr==0.4.15
2
+ pypdfium2==4.30.0
3
+ openai==1.35.13