62 lines
1.8 KiB
Python
62 lines
1.8 KiB
Python
import os
|
|
import cv2
|
|
import easyocr
|
|
import numpy as np
|
|
from core.response import SuccessResponse
|
|
from fastapi import APIRouter, UploadFile, Form
|
|
from schemas.orc_result import ResultInfo, ResultMain
|
|
|
|
|
|
app = APIRouter()
|
|
|
|
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
model_dir = os.path.join(BASE_DIR, 'model')
|
|
|
|
reader = easyocr.Reader(
|
|
['ch_sim', 'en'],
|
|
model_storage_directory=model_dir,
|
|
download_enabled=False,
|
|
gpu=True)
|
|
|
|
|
|
def preprocess_img(image):
|
|
# 转为灰度图
|
|
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
# 二值化
|
|
_, binary = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
|
|
# 去噪声
|
|
denoised = cv2.medianBlur(binary, 3)
|
|
return denoised
|
|
|
|
|
|
@app.post("/upload")
|
|
async def orc(files: list[UploadFile] = Form(...)):
|
|
result = []
|
|
for file in files:
|
|
main = ResultMain()
|
|
main.file_name = file.filename
|
|
# 读取上传的文件内容
|
|
image_data = await file.read()
|
|
|
|
# 2. 将字节数据转换为 numpy 数组
|
|
np_array = np.frombuffer(image_data, np.uint8)
|
|
|
|
# 3. 使用 cv2.imdecode 解码为 OpenCV 图像
|
|
image = cv2.imdecode(np_array, cv2.IMREAD_COLOR)
|
|
new_image = preprocess_img(image)
|
|
datas = reader.readtext(np.array(new_image))
|
|
infos = []
|
|
for data in datas:
|
|
info = ResultInfo()
|
|
bounding_boxs = data[0]
|
|
left_up = bounding_boxs[0]
|
|
right_down = bounding_boxs[2]
|
|
info.bounding_box_left_up = [int(left_up[0]), int(left_up[1])]
|
|
info.bounding_box_right_down = [int(right_down[0]), int(right_down[1])]
|
|
info.text = data[1]
|
|
info.confidence = round(data[2], 4)
|
|
infos.append(info)
|
|
main.infos = infos
|
|
result.append(main.model_dump())
|
|
return SuccessResponse(data=result)
|