import gradio as gr

from askui_ml_helper.utils.pta_text import PtaTextInference

def main():
    inference = PtaTextInference("model/pta-text-v0.1.pt")
    # Gradio Interface
    iface = gr.Interface(
        fn=inference.process_image_and_draw_circle,
        inputs=[gr.Image(type="pil", label = "Upload Image"),
                gr.Textbox(label = "Prompt", placeholder='Enter a prompt to find in the image', value='click on the text "Login"'),
                gr.Slider(minimum=5, maximum=20, value=10, label="Prediction Dot Size", step=5),
                gr.Checkbox(label="Resize (Recommended)", value=True,  info="Resize the image to 1920x1080 for better results?")],
        outputs=gr.Image(type="pil"),
        title="PTA-Text: A **Text-only** Click Model",
        description="Based on PTA-text model ([Huggingface model card](https://huggingface.co/AskUI/pta-text-0.1)) from [askui](https://www.askui.com/).PTA is a small part of our bigger model - PTA-1 (Prompt-To-Automation). \nInstructions: Upload a UI image and enter a text finding prompt to see the model's prediction. \nLimitations: Currently, limited to single screen images.",
        examples=[["./sample.png", 'click on the text "Like"', 10, True],
                ["./sample.png", 'click on the text "Messaging" on the upper part of the screen', 15, True],
                ["./sample.png", 'click on the text "Retry for free!"', 15, True]]
    )


    iface.launch(server_name="0.0.0.0")



if __name__ == "__main__":
    main()