import gradio as gr
import os
import cv2
import numpy as np
from PIL import Image
from paddleocr import PPStructure, save_structure_res, draw_structure_result
from paddleocr.ppstructure.recovery.recovery_to_doc import sorted_layout_boxes, convert_info_docx

# Chinese image
# table_engine = PPStructure(recovery=True)

# English image
# table_engine = PPStructure(recovery=True, lang='en')

examples = [['0.png']]

def find_layout(image, mode):
    save_folder = './output'
    # img = cv2.imread(image)
    img = np.array(image)

    print("The variable image is of type:", type(image))
    print("The variable img is of type:", type(img))

    if mode is not None:
        mode = mode[0]
    print("Mode: ", mode)
    if mode=="LayoutAnalysis":
        table_engine = PPStructure(table=False, ocr=False, show_log=True)
    elif mode=="TableRecognition":
        table_engine = PPStructure(layout=False, show_log=True)
    else:
        table_engine = PPStructure(recovery=True)
    
    result = table_engine(img)
    save_structure_res(result, save_folder, os.path.basename("result").split('.')[0])

    if mode=="LayoutReccovery":
        h, w, _ = img.shape
        res = sorted_layout_boxes(result, w)
    
    final_text = []
    for line in result:
        line.pop('img')
        print(line)
        final_text.append(line)

    # convert_info_docx(img, res, save_folder, os.path.basename("result").split('.')[0])

    font_path = 'simfang.ttf' # PaddleOCR下提供字体包

    # image = Image.open(img_path).convert('RGB')
    image = image.convert('RGB')
    # img = img.convert('RGB')
    im_show = draw_structure_result(image, result, font_path=font_path)
    im_show = Image.fromarray(im_show)
    # im_show.save('result.jpg')

    return final_text, im_show

iface = gr.Interface(fn=find_layout, inputs=[
        gr.Image(type="pil"),
        gr.CheckboxGroup(["LayoutAnalysis", "TableRecognition", "LayoutRecovery"], label="mode"),
    ], outputs=["text", "image"], examples=examples)
iface.launch()