from apis.layoutlm import LayoutLM
import pandas as pd
import gradio as gr
import os

layoutlm = None

def auth(username, password):
    u = os.environ.get('USERNAME')
    p = os.environ.get('PASSWORD')
    return (username == u and password == p)

def inference(img) -> pd.DataFrame:
    return layoutlm.inference(img)

if __name__ == "__main__":

    try:
        layoutlm = LayoutLM()
        layoutlm.set_model(layoutlm.default_model)

        demo = gr.Interface(
            inference,
            gr.Image(type="pil"),
            gr.Dataframe(
                headers=['Data', 'Value'],
                datatype=['str', 'str'],
                row_count=8,
                col_count=(2, 'fixed'), 
                interactive=False
            )
        )

        demo.launch(auth=auth)

    except Exception as e:
        print(str(e))