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