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)