105 lines
3.2 KiB
Python
105 lines
3.2 KiB
Python
from fastapi import FastAPI, Request, UploadFile, File, Form
|
|
from fastapi.templating import Jinja2Templates
|
|
from fastapi.staticfiles import StaticFiles
|
|
from fastapi.responses import HTMLResponse, JSONResponse
|
|
import shutil
|
|
import os
|
|
import tempfile
|
|
import asyncio
|
|
from typing import List, Optional
|
|
|
|
from LM import run_pdf, run_images
|
|
|
|
|
|
try:
|
|
with open("InvoicePrompt", 'r', encoding='utf-8') as f:
|
|
SYS_INVOICE_PROMPT = f.read()
|
|
except FileNotFoundError:
|
|
SYS_INVOICE_PROMPT = "请识别发票内容并返回JSON格式数据。"
|
|
print("Warning: InvoicePrompt file not found.")
|
|
|
|
try:
|
|
with open("OCR_MD", 'r', encoding='utf-8') as f:
|
|
SYS_OCR_PROMPT = f.read()
|
|
except FileNotFoundError:
|
|
SYS_OCR_PROMPT = "请识别图片中的文字并转换为Markdown格式。"
|
|
print("Warning: OCR_MD file not found.")
|
|
|
|
app = FastAPI()
|
|
|
|
templates = Jinja2Templates(directory="templates")
|
|
|
|
app.mount("/static", StaticFiles(directory="static"), name="static")
|
|
|
|
|
|
@app.get("/", response_class=HTMLResponse)
|
|
async def read_root(request: Request):
|
|
|
|
return templates.TemplateResponse("index2.html", {"request": request})
|
|
|
|
|
|
@app.post("/api/recognize")
|
|
async def recognize_files(
|
|
files: List[UploadFile] = File(...),
|
|
mode: str = Form(...),
|
|
custom_prompt: str = Form(None)
|
|
):
|
|
results = []
|
|
|
|
final_prompt = ""
|
|
if mode == 'invoice':
|
|
final_prompt = SYS_INVOICE_PROMPT
|
|
print("Using System Invoice Prompt")
|
|
elif mode == 'ocr':
|
|
final_prompt = SYS_OCR_PROMPT
|
|
print("Using System OCR Prompt")
|
|
elif mode == 'custom':
|
|
final_prompt = custom_prompt if custom_prompt else ""
|
|
print(f"Using Custom Prompt (Length: {len(final_prompt)})")
|
|
else:
|
|
final_prompt = SYS_INVOICE_PROMPT
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
tasks = []
|
|
|
|
for file in files:
|
|
file_path = os.path.join(temp_dir, file.filename)
|
|
|
|
with open(file_path, "wb") as buffer:
|
|
shutil.copyfileobj(file.file, buffer)
|
|
|
|
original_name = file.filename
|
|
ext = os.path.splitext(original_name)[1].lower()
|
|
|
|
if ext == ".pdf":
|
|
tasks.append(run_pdf(original_name, file_path, prompt=final_prompt))
|
|
|
|
elif ext in [".jpg", ".jpeg", ".png", ".bmp", ".webp"]:
|
|
tasks.append(run_images([original_name], [file_path], prompt=final_prompt))
|
|
|
|
else:
|
|
results.append({
|
|
"image_name": original_name,
|
|
"result": "Unsupported file type",
|
|
"state": "error"
|
|
})
|
|
|
|
if tasks:
|
|
processed_results = await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
for res in processed_results:
|
|
if isinstance(res, Exception):
|
|
print(f"Internal Error: {str(res)}")
|
|
results.append({
|
|
"state": "error",
|
|
"result": f"Internal execution error: {str(res)}"
|
|
})
|
|
else:
|
|
results.append(res)
|
|
|
|
return JSONResponse(content={"data": results})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
uvicorn.run(app, host="0.0.0.0", port=9420) |