from src.rep_api import virtual_try_on
import gradio as gr 


def create_virtualtryon_tab():
    with gr.TabItem("Virtual Try-on",id="virtual_tryon"):
        with gr.Row():
            with gr.Column():
                human_img =  gr.Image(label="Human Image")
                crop = gr.Checkbox(label="Crop", value=False)
                seed = gr.Number(label="Seed", value=42)
                steps = gr.Number(label="Steps", value=30)
                with gr.Accordion("Example People", open=False):
                    human_examples = gr.Examples(
                        examples=[
                            ["Test_images/Woman_1.png"],
                            ["Test_images/person_1.jpg"],
                            ["Test_images/Woman_2.png"],
                            ["Test_images/person_2.jpg"],
                            ["Test_images/Woman_3.png"],
                            ["Test_images/man_1.png"],
                            ["Test_images/Woman_4.png"],
                             ["Test_images/Woman_5.png"],
                            ["Test_images/anime_woman_1.png"],
                             ["Test_images/anime_woman_2.png"],
                            ["Test_images/anime_woman_3.png"],
                            ["Test_images/Jump.png"],
                            ["Test_images/Walk_1.png"],
                            ["Test_images/Walk_2.png"],
                            ["Test_images/Walk_girl_1.png"],
                            ["Test_images/Walk_girl_2.png"]
                        ],
                        inputs=[human_img],
                         examples_per_page=24
                    )
            with gr.Column():
                garm_img = gr.Image(label="Garment Image")
                category = gr.Dropdown(["upper_body", "lower_body", "dresses"], label="Category", value="upper_body")
                garment_des = gr.Textbox(label="Garment Description")
                with gr.Accordion("Example Clothes", open=False):
                    cloth_examples = gr.Examples(
                        examples=[
                            ["Test_images/pink_jumper.png","pink jumper","upper_body"],
                            ["Test_images/Suit_2.png","Suit","upper_body"],
                            ["Test_images/Jacket_1.png","Colorfull leather jacket","upper_body"],
                            ["Test_images/Jacket_2.png","White jacket","upper_body"],
                            ["Test_images/Shirt_1.png","Shirt","upper_body"],
                            ["Test_images/dress_2.png","dress","dresses"],
                            ["Test_images/dress_1.png","dress","dresses"],
                            ["Test_images/dress_3.png","pink dress","dresses"],
                            ["Test_images/dress_4.png","red dress","dresses"],
                             ["Test_images/dress_5.png","golden dress","dresses"],
                              ["Test_images/Shoes_1.png","Shoes","lower_body"],
                              ["Test_images/Shoes_2.png","Shoes","lower_body"],
                              ["Test_images/Nike_RTFKT_Dunk_Void_1.jpg","lower_body"]

                        ],
                        inputs=[garm_img,garment_des,category],
                        examples_per_page=4
                    )

        with gr.Row():
            tryon_btn = gr.Button("Try-on")
        with gr.Row():
            tryon_out = gr.Image(label="Output", type="filepath")
            
       
    tryon_btn.click(virtual_try_on, inputs=[crop, seed, steps, category,  garm_img, human_img,garment_des],outputs=tryon_out)

    return human_img,tryon_out