from core.response import SuccessResponse
from fastapi import APIRouter, UploadFile, Form
from schemas.orc_result import ResultInfo, ResultMain

from paddleocr import PaddleOCR
import numpy as np
import cv2


app = APIRouter()


paddle_ocr = PaddleOCR(lang='ch')


def enhance_image(img):
    """增强亮部和暗部，提升数字清晰度"""
    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
    cl = clahe.apply(l)
    enhanced = cv2.merge((cl, a, b))
    return cv2.cvtColor(enhanced, cv2.COLOR_LAB2BGR)


@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()
        # 从字节数据读取图像
        np_array = np.frombuffer(image_data, np.uint8)
        img = cv2.imdecode(np_array, cv2.IMREAD_COLOR)
        # 图像增强
        img = enhance_image(img)
        # 4. 调用 OCR 模型进行识别
        datas = paddle_ocr.ocr(img, cls=False)
        infos = []
        if datas:
            for data in datas[0]:
                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][0]
                info.confidence = round(data[1][1], 4)
                infos.append(info)
            main.infos = infos
            result.append(main.model_dump())
    return SuccessResponse(data=result)