377 lines
12 KiB
Python
377 lines
12 KiB
Python
import uuid
|
||
|
||
from fastapi import FastAPI, File, UploadFile, HTTPException, Request
|
||
import ssl
|
||
|
||
from uvicorn import Config, Server
|
||
from fastapi.responses import JSONResponse
|
||
from fastapi.encoders import jsonable_encoder
|
||
import requests
|
||
import uvicorn
|
||
import config
|
||
from pydantic import BaseModel
|
||
from typing import Union, List
|
||
import fitz
|
||
from PIL import Image
|
||
import json
|
||
import re
|
||
import os
|
||
from openai import OpenAI
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
from fastapi.staticfiles import StaticFiles
|
||
|
||
from classify import predict_folder
|
||
from ocr import get_ocr_list, get_ocr_image_list, get_ocr
|
||
from torchvision import transforms, models
|
||
|
||
|
||
app = FastAPI()
|
||
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=["*"],
|
||
allow_credentials=True,
|
||
allow_methods=["*"],
|
||
allow_headers=["*"]
|
||
)
|
||
|
||
app.mount("/data", StaticFiles(directory="data"), name="data")
|
||
|
||
# 确保上传目录存在
|
||
current_directory = os.getcwd()
|
||
upload_dir = os.path.join(current_directory, "data\\")
|
||
images_dir = upload_dir + "img\\"
|
||
pdfs_dir = upload_dir + "content\\"
|
||
os.makedirs(images_dir, exist_ok=True)
|
||
os.makedirs(pdfs_dir, exist_ok=True)
|
||
|
||
|
||
def run_conv(text_prompt, ocr_text):
|
||
print("begin deepseek ask")
|
||
messages = [
|
||
{'role': 'user', 'content': str(ocr_text)},
|
||
{'role': 'assistant', 'content': text_prompt}
|
||
]
|
||
|
||
# client = OpenAI(
|
||
# api_key="sk-0ffaa2ae7c5c499aa7fd03e646b6717a",
|
||
# base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
|
||
# )
|
||
client = OpenAI(api_key='ollama',
|
||
base_url="http://localhost:11434/v1")
|
||
response = client.chat.completions.create(
|
||
# model="qwen-plus",
|
||
model = "deepseek-r1",
|
||
messages=messages
|
||
)
|
||
final_response = response.choices[0].message.content
|
||
match = re.search(r"```json(.*?)```", final_response, re.DOTALL)
|
||
res = match.group(1)
|
||
print("llm result: ", res)
|
||
|
||
return res
|
||
|
||
|
||
# 发票
|
||
@app.post("/invoice/")
|
||
async def invoice(files: List[UploadFile] = File(...)):
|
||
image_url = []
|
||
for file in files:
|
||
if file.content_type.startswith('image/'):
|
||
# 保存图片文件
|
||
file_location = images_dir + file.filename
|
||
else:
|
||
raise HTTPException(status_code=400, detail="Unsupported file type")
|
||
|
||
with open(file_location, "wb+") as file_object:
|
||
file_object.write(await file.read())
|
||
|
||
image_url.append(file_location)
|
||
print("image_path: ", image_url)
|
||
|
||
try:
|
||
text_prompt = config.invoice
|
||
|
||
ocr_text = get_ocr_list(image_url) # 进行OCR，返回 String
|
||
output = run_conv(text_prompt, ocr_text) # 调用大模型
|
||
new_url = get_ocr_image_list(image_url, output)
|
||
|
||
return JSONResponse(content={"text": output,"image_url": new_url, "class": "发票"})
|
||
|
||
except Exception as e:
|
||
return JSONResponse(status_code=400, content={"error": str(e)})
|
||
|
||
# 申请表
|
||
@app.post("/application/")
|
||
async def application(files: List[UploadFile] = File(...)):
|
||
image_url = []
|
||
for file in files:
|
||
if file.content_type.startswith('image/'):
|
||
# 保存图片文件
|
||
file_location = images_dir + file.filename
|
||
else:
|
||
raise HTTPException(status_code=400, detail="Unsupported file type")
|
||
|
||
with open(file_location, "wb+") as file_object:
|
||
file_object.write(await file.read())
|
||
|
||
image_url.append(file_location)
|
||
|
||
try:
|
||
text_prompt = config.application
|
||
ocr_text = get_ocr_list(image_url) # 进行OCR，返回String
|
||
output = run_conv(text_prompt, ocr_text) # 调用大模型
|
||
|
||
return JSONResponse(content={"text": output, "image_url":image_url, "class":"申请单"})
|
||
except Exception as e:
|
||
return JSONResponse(status_code=400, content={"error": str(e)})
|
||
|
||
# 确认表
|
||
@app.post("/confirmation/")
|
||
async def confirmation(files: List[UploadFile] = File(...)):
|
||
image_url = []
|
||
|
||
for file in files:
|
||
if file.content_type.startswith('image/'):
|
||
# 保存图片文件
|
||
file_location = images_dir + file.filename
|
||
else:
|
||
raise HTTPException(status_code=400, detail="Unsupported file type")
|
||
|
||
with open(file_location, "wb+") as file_object:
|
||
file_object.write(await file.read())
|
||
|
||
image_url.append(file_location)
|
||
|
||
try:
|
||
text_prompt = config.confirmation
|
||
ocr_text = get_ocr_list(image_url)
|
||
output = run_conv(text_prompt, ocr_text)
|
||
|
||
return JSONResponse(content={"text": output, "image_url":image_url, "class":"确认表"})
|
||
except Exception as e:
|
||
return JSONResponse(status_code=400, content={"error": str(e)})
|
||
|
||
# 合同
|
||
@app.post("/contract/")
|
||
async def contract(files: UploadFile = File(...)):
|
||
if files.content_type != 'application/pdf':
|
||
raise HTTPException(status_code=400, detail="文件类型不符")
|
||
|
||
random_filename = str(uuid.uuid4())
|
||
file_location = pdfs_dir + f"{random_filename}.pdf"
|
||
|
||
with open(file_location, "wb+") as file_object:
|
||
file_object.write(await files.read())
|
||
print("合同上传位置：", file_location)
|
||
|
||
try:
|
||
img_url, ocr_text = pagehome_llm(file_location)
|
||
text_prompt = config.contract
|
||
print("contract img_url: ", img_url)
|
||
|
||
output = run_conv(text_prompt, ocr_text)
|
||
new_url = get_ocr_image_list(img_url, output)
|
||
|
||
return JSONResponse(content={"text": output, "image_url": new_url, "class": "合同表"})
|
||
except Exception as e:
|
||
return JSONResponse(status_code=400, content={"error": str(e)})
|
||
|
||
|
||
@app.post("/contractPayment/")
|
||
async def contract_amount(files: UploadFile = File(...)):
|
||
if files.content_type != 'application/pdf':
|
||
raise HTTPException(status_code=400, detail="Unsupported file type")
|
||
|
||
random_filename = str(uuid.uuid4())
|
||
file_location = pdfs_dir + f"{random_filename}.pdf"
|
||
|
||
with open(file_location, "wb+") as file_object:
|
||
file_object.write(await files.read())
|
||
print("合同上传位置：", file_location)
|
||
|
||
try:
|
||
json_data = """
|
||
{
|
||
"page":"",
|
||
"boolean":""
|
||
}
|
||
"""
|
||
result = json.loads(json_data)
|
||
|
||
# 打开PDF文件
|
||
pdf_document = fitz.open(file_location)
|
||
for page_num in range(len(pdf_document)):
|
||
image_filename = pdf_2_images(pdf_document, images_dir, page_num)
|
||
|
||
ocr_text = get_ocr(image_filename)
|
||
if "合同价款暂定为人民币含税价小写" in str(ocr_text):
|
||
page = re.search(r'(\d+)\.png', image_filename).group(1)
|
||
result["page"] = page
|
||
if "合同价款结算按第" in str(ocr_text):
|
||
pattern = r"合同价款结算按第(\d+)_种方式"
|
||
match = re.search(pattern, str(ocr_text))
|
||
if match.group(1) == "2":
|
||
result["boolean"] = "是"
|
||
break
|
||
|
||
if result["page"] == "":
|
||
result["page"] = "0"
|
||
if result["boolean"] == "":
|
||
result["boolean"] = "否"
|
||
return JSONResponse(content={"text": result, "pdf_url": file_location, "class": "合同抽取"})
|
||
|
||
except Exception as e:
|
||
return JSONResponse(status_code=400, content={"error": str(e)})
|
||
|
||
|
||
def pdf_2_images(pdf_document, output_folder, page_num):
|
||
# 获取页面
|
||
page = pdf_document.load_page(page_num)
|
||
|
||
# 将页面转换成图像
|
||
pix = page.get_pixmap()
|
||
|
||
# 使用PIL保存图像
|
||
img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
|
||
|
||
match = re.search(r'\\([^\\]+)\.pdf$', pdf_document.name)
|
||
contract_name = match.group(1)
|
||
image_filename = "{}{}page_{}.png".format(output_folder, contract_name, page_num + 1)
|
||
img.save(image_filename, "PNG")
|
||
print("合同分页保存位置：", image_filename)
|
||
return image_filename
|
||
|
||
def pagehome_llm(pdf_path):
|
||
result = []
|
||
ocr_text = []
|
||
pdf_document = fitz.open(pdf_path)
|
||
|
||
for page_num in range(len(pdf_document)):
|
||
image_filename = pdf_2_images(pdf_document, images_dir, page_num)
|
||
ocr_text.append(get_ocr(image_filename))
|
||
res = "甲方" and "乙方" and "合同编号" and "签订地点" in str(ocr_text)
|
||
|
||
if res:
|
||
print("合同识别起始页：", page_num)
|
||
result.append(image_filename)
|
||
|
||
image_filename = pdf_2_images(pdf_document, images_dir, page_num + 1)
|
||
ocr_text.append(get_ocr(image_filename))
|
||
result.append(image_filename)
|
||
|
||
image_filename = pdf_2_images(pdf_document, images_dir, page_num + 2)
|
||
ocr_text.append(get_ocr(image_filename))
|
||
result.append(image_filename)
|
||
|
||
return result, ocr_text
|
||
|
||
return result, ocr_text
|
||
|
||
|
||
|
||
|
||
|
||
def find_payment_page(pdf_path, output_folder):
|
||
# text_prompt = config.findPage
|
||
|
||
pdf_document = fitz.open(pdf_path)
|
||
for page_num in range(len(pdf_document)):
|
||
image_filename = pdf_2_images(pdf_document, images_dir, page_num)
|
||
ocr_text = get_ocr(image_filename)
|
||
if "合同价款暂定为人民币含税价小写" in str(ocr_text):
|
||
result = re.search(r'(\d+)\.png', image_filename).group(1)
|
||
print("位置为第" + result + "页")
|
||
return result
|
||
|
||
def extract_images(image_index, pdf_path,output_folder):
|
||
# images_pdf = convert_from_path(pdf_path)
|
||
pdf_document = fitz.open(pdf_path)
|
||
index = int(image_index)
|
||
i = 0
|
||
ocr_texts = []
|
||
text_prompt = config.extract
|
||
if len(pdf_document) >= index:
|
||
while i < 2:
|
||
image_filename = pdf_2_images(pdf_document, images_dir, index-1)
|
||
ocr_text = get_ocr(image_filename)
|
||
ocr_texts.append(ocr_text)
|
||
i += 1
|
||
index += 1
|
||
else:
|
||
print("超出PDF文件长度")
|
||
|
||
output = run_conv(text_prompt, ocr_texts)
|
||
output = re.sub(r'[\x00-\x1F\x7F-\x9F]', '', output)
|
||
json_data = json.loads(output)
|
||
if json_data["result"]:
|
||
result = json_data["result"]
|
||
return JSONResponse(content=result)
|
||
|
||
@app.post("/findAndExtract/")
|
||
async def findAndExtract(request:Request):
|
||
|
||
body = await request.json()
|
||
file_path = body.get("file_path")
|
||
|
||
# 检查文件夹路径是否存在
|
||
if not os.path.exists(file_path):
|
||
return JSONResponse(content={"error": "File path does not exist"}, status_code=400)
|
||
|
||
index = find_payment_page(file_path, images_dir) # 遍历到第几页
|
||
result = extract_images(index, file_path, images_dir)
|
||
json_result = jsonable_encoder(result)
|
||
return JSONResponse(content=json_result)
|
||
|
||
|
||
@app.post('/classify')
|
||
async def get_images(request: Request):
|
||
body = await request.json()
|
||
folder_path = body.get("file_path")
|
||
|
||
# 检查文件夹路径是否存在
|
||
if not os.path.exists(folder_path):
|
||
return JSONResponse(content={"error": "File path does not exist"}, status_code=400)
|
||
|
||
class_names = ['发票', '确认表', '申请表', '验收证书', '其他']
|
||
# 训练权重路径
|
||
weights_path = "../Models/checkpoints/model.pth"
|
||
|
||
# 对图像进行变换
|
||
transform = transforms.Compose([
|
||
transforms.Resize((224, 224)),
|
||
transforms.RandomHorizontalFlip(),
|
||
transforms.RandomRotation(10),
|
||
transforms.ColorJitter(brightness=0.2, contrast=0.2),
|
||
transforms.ToTensor(),
|
||
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
|
||
])
|
||
|
||
result = predict_folder(folder_path, weights_path, transform, class_names)
|
||
json_result = jsonable_encoder(result)
|
||
return JSONResponse(content=json_result)
|
||
|
||
if __name__ == "__main__":
|
||
# ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
|
||
# ssl_context.load_cert_chain(certfile="/home/oem/llm/certificate.crt",keyfile="/home/oem/llm/private.key")
|
||
# config = Config(app=app, host="0.0.0.0", port=8000, ssl_context=ssl_context)
|
||
# uvicorn.run(app, host="0.0.0.0", port=8000,ssl=ssl_context)
|
||
|
||
uvicorn.run(app, host="0.0.0.0", port=8000)
|
||
|
||
# weights_path = "../Models/checkpoints/model_epoch_1.pth"
|
||
#
|
||
# # 对图像进行变换
|
||
# transform = transforms.Compose([
|
||
# transforms.Resize((224, 224)),
|
||
# transforms.RandomHorizontalFlip(),
|
||
# transforms.RandomRotation(10),
|
||
# transforms.ColorJitter(brightness=0.2, contrast=0.2),
|
||
# transforms.ToTensor(),
|
||
# transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
|
||
# ])
|
||
# class_names = ['发票', '确定表', '申请表']
|
||
# print(predict_folder("data/img/分类", weights_path, transform, class_names))
|
||
# app.run(host='0.0.0.0', port=5001)
|
||
|