from rapidocr import RapidOCR

from rapid_table import ModelType, RapidTable, RapidTableInput

ocr_engine = RapidOCR()

ROOT_URL = "models/"
model_path = {
    "encoder": f"{ROOT_URL}/unitable/encoder.pth",
    "decoder": f"{ROOT_URL}/unitable/decoder.pth",
    "vocab": f"{ROOT_URL}/unitable/vocab.json",
}

input_args = RapidTableInput(
     model_type=ModelType.UNITABLE,
     model_dir_or_path=model_path,
     engine_cfg={"use_cuda": True, "cuda_ep_cfg.gpu_id": 1},
)

table_engine = RapidTable(input_args)


@app.post("/uni_table")
def uni_table(
     image_data: ImageData,
) -> str:
    image_bytes = base64.b64decode(image_data.image)
    ori_ocr_res  = ocr_engine(image_bytes)
    ocr_results = [ori_ocr_res.boxes, ori_ocr_res.txts, ori_ocr_res.scores]
    table_results = table_engine(image_bytes , ocr_results)
    table_html_str, table_cell_bboxes = table_results.pred_html, table_results.cell_bboxes
    return table_html_str

def start_server(
        host='0.0.0.0', port=0000, reload=False, **kwargs
):
    import uvicorn

    uvicorn.run(app, host=host, port=port, reload=reload, **kwargs)