patharanor/invex

1from apis.layoutlm import LayoutLM2from apis.qc3.host_list import HostList3from PIL import Image4import pandas as pd5import gradio as gr6import os7 8layoutlm = None9hl = HostList(is_debug=True)10 11def auth(username, password):12    u = os.environ.get('USERNAME')13    p = os.environ.get('PASSWORD')14    return (username == u and password == p)15 16def inference(img) -> pd.DataFrame:17    return layoutlm.inference(img)18 19def filter_green_out(img: Image):20    image_data = img.load()21    height,width = img.size22    for loop1 in range(height):23        for loop2 in range(width):24            (r,g,b) = image_data[loop1,loop2]25            if g < 70 and r < 70 and b < 70:26                (r,g,b) = (0,0,0)27            else:28                (r,g,b) = (255,255,255)29            image_data[loop1,loop2] = r,g,b30    31    img.save('./temp.jpg')32    return img33 34def ask(img: Image, question, top_k, max_answer_len, chk_is_remove_green) -> str:35    if chk_is_remove_green:36        img = filter_green_out(img.copy())37 38    return layoutlm.answer_the_question_without_filter(39        img, 40        question, 41        top_k=top_k, 42        max_answer_len=max_answer_len, 43        is_debug=True)44 45if __name__ == '__main__':46 47    try:48        layoutlm = LayoutLM()49        layoutlm.set_model(layoutlm.default_model)50 51        with gr.Blocks() as demo:52 53            with gr.Tab('List'):54                with gr.Row():55                    with gr.Column():56                        list_inp_img = gr.Image(type="pil")57                        gr.Examples(58                            [['./examples/host-list1.JPG'], ['./examples/host-list2.JPG', './examples/host-list3.JPG']], 59                            list_inp_img60                        )61 62                    with gr.Column():63                        list_out_txt = gr.Textbox(label='Answer', interactive=False)64 65                        list_btn_ask = gr.Button('Ask me')66                        list_btn_ask.click(hl.process_image, [67                            list_inp_img68                        ], list_out_txt)69 70            with gr.Tab('Layout'):71                with gr.Row():72                    inp_img = gr.Image(type='pil')73                74                    with gr.Column():75                        out = gr.Dataframe(76                            headers=['Data', 'Value'],77                            datatype=['str', 'str'],78                            row_count=8,79                            col_count=(2, 'fixed'), 80                            interactive=False81                        )82 83                        txt_custom_question = gr.Textbox(label='Your question')84                        sld_max_answer = gr.Slider(1, 10, value=1, step=1, label="Max answer", info="Top-K between 1 and 10")85                        sld_max_answer_len = gr.Slider(1, 200, value=15, step=1, label="Max answer length", info="Length between 15 and 200")86                        chk_is_remove_green = gr.Checkbox(label="Remove green", info="Do you need clean context?")87                        btn_ask = gr.Button('Ask me')88                        txt_out_answer = gr.Textbox(label='Answer', interactive=False)89 90                        # event91                        inp_img.change(inference, inp_img, out)92                        btn_ask.click(ask, [93                            inp_img, 94                            txt_custom_question, 95                            sld_max_answer, 96                            sld_max_answer_len,97                            chk_is_remove_green98                        ], txt_out_answer)99 100        #demo.launch(auth=auth)101        demo.launch()102    except Exception as e:103        print(str(e))