121 lines
5.5 KiB
Python
121 lines
5.5 KiB
Python
import contextlib
|
|
import json
|
|
import os
|
|
import shutil
|
|
import sys
|
|
import traceback
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
from domain import auto_match, enrich, extract
|
|
from exports import export_expense, export_ppt, export_travel
|
|
|
|
PROTOCOL = sys.stdout
|
|
SUPPORTED = {'.jpg', '.jpeg', '.png', '.webp', '.bmp', '.gif', '.pdf'}
|
|
|
|
|
|
def emit(value):
|
|
PROTOCOL.write(json.dumps(value, ensure_ascii=False) + '\n')
|
|
PROTOCOL.flush()
|
|
|
|
|
|
def scan(request):
|
|
import onnxruntime
|
|
import pypdfium2 as pdfium
|
|
from PIL import Image
|
|
from rapidocr_onnxruntime import RapidOCR
|
|
root = Path(request['path'])
|
|
destination = Path(request['workspacePath'])
|
|
if not root.is_dir():
|
|
raise ValueError('请选择有效的报账材料文件夹')
|
|
destination.mkdir(parents=True, exist_ok=True)
|
|
state = dict(rootPath=request.get('folderName', root.name), invoices=[], payments=[], photos=[], matches=[], warnings=[], directoryPaymentTotal=0)
|
|
files = []
|
|
for source in sorted(root.rglob('*')):
|
|
if source.is_file() and source.suffix.lower() in SUPPORTED:
|
|
parts = source.relative_to(root).parts[:-1]
|
|
kind = next((kind for label, kind in [('发票', 'invoice'), ('付款截图', 'payment'), ('实物照片', 'photo')] if label in parts), None)
|
|
if kind:
|
|
files.append((source, kind))
|
|
if not files:
|
|
raise ValueError('目录中没有可处理材料，请确认包含“发票”“付款截图”子文件夹')
|
|
emit(dict(event='progress', progress=0, message='正在加载本地 RapidOCR 模型'))
|
|
onnxruntime.disable_telemetry_events()
|
|
recognizer = RapidOCR()
|
|
for index, (source, kind) in enumerate(files):
|
|
identifier = str(uuid.uuid4())
|
|
stored = destination / (identifier + source.suffix.lower())
|
|
shutil.copyfile(source, stored)
|
|
item = dict(id=identifier, name=source.name, path=str(stored), type=kind, size=stored.stat().st_size, matched=False,
|
|
previewPath='', ocr=extract([]))
|
|
item['ocr']['status'] = 'pending'
|
|
try:
|
|
if source.suffix.lower() == '.pdf':
|
|
with contextlib.closing(pdfium.PdfDocument(str(stored))) as document:
|
|
if not len(document):
|
|
raise ValueError('PDF 没有可预览页面')
|
|
with contextlib.closing(document[0]) as page:
|
|
with contextlib.closing(page.render(scale=140 / 72)) as bitmap:
|
|
image = bitmap.to_pil().convert('RGB')
|
|
else:
|
|
with Image.open(stored) as original:
|
|
image = original.convert('RGB')
|
|
preview = destination / (identifier + '-preview.png')
|
|
image.save(preview, format='PNG')
|
|
item['previewPath'] = str(preview)
|
|
if kind != 'photo':
|
|
result, _ = recognizer(str(preview if source.suffix.lower() == '.pdf' else stored))
|
|
item['ocr'] = extract([str(entry[1]).strip() for entry in result or [] if len(entry) >= 3 and str(entry[1]).strip()])
|
|
except Exception as error:
|
|
item['ocr']['status'] = 'failed'
|
|
item['ocr']['message'] = str(error)
|
|
state['warnings'].append(source.name + '：识别或预览失败，可人工配对')
|
|
state[{'invoice': 'invoices', 'payment': 'payments', 'photo': 'photos'}[kind]].append(item)
|
|
emit(dict(event='progress', progress=(index + 1) / len(files) * .95, message=f'已处理 {index + 1}/{len(files)} · {source.name}'))
|
|
for field, label in [('invoices', '发票'), ('payments', '付款截图'), ('photos', '实物照片')]:
|
|
if not state[field]:
|
|
state['warnings'].append('未找到“' + label + '”文件夹或支持的文件')
|
|
emit(dict(event='progress', progress=.97, message='正在执行自动配对与分类'))
|
|
return enrich(auto_match(state))
|
|
|
|
|
|
def dispatch(request):
|
|
operation = request['operation']
|
|
if operation == 'scan':
|
|
return scan(request)
|
|
state = request['state']
|
|
if operation == 'refresh':
|
|
return enrich(state)
|
|
if not (state.get('approvedRecords') if operation == 'approved-ppt' else state.get('matches')):
|
|
raise ValueError('至少完成一组核对后才能导出')
|
|
destination = Path(request['destination'])
|
|
temporary = destination.with_name('.' + str(uuid.uuid4()) + destination.suffix)
|
|
metadata = {}
|
|
try:
|
|
if operation in ('ppt', 'approved-ppt'):
|
|
export_ppt(state, temporary, request.get('classified', True))
|
|
elif operation == 'travel':
|
|
export_travel(state, temporary)
|
|
elif operation == 'expense':
|
|
template = Path(__file__).parent / 'personal-expense-template.xlsx'
|
|
metadata = export_expense(state, temporary, template, request.get('purposes', {}), request.get('signatures', []),
|
|
request.get('payee', {}), request.get('categoryPurposes'))
|
|
else:
|
|
raise ValueError('未知操作：' + operation)
|
|
os.replace(temporary, destination)
|
|
finally:
|
|
temporary.unlink(missing_ok=True)
|
|
return dict(destination=str(destination), **metadata)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
try:
|
|
request = json.loads(sys.stdin.readline())
|
|
with contextlib.redirect_stdout(sys.stderr):
|
|
result = dispatch(request)
|
|
emit(dict(event='result', result=result))
|
|
except Exception as error:
|
|
traceback.print_exc(file=sys.stderr)
|
|
emit(dict(event='error', message=str(error)))
|
|
sys.exit(1)
|