import torch
from ram import get_transform, inference_ram, inference_tag2text
from ram.models import ram, tag2text

ram_checkpoint = "./ram_swin_large_14m.pth"
tag2text_checkpoint = "./tag2text_swin_14m.pth"
image_size = 384
device = "cuda" if torch.cuda.is_available() else "cpu"


@torch.no_grad()
def inference(raw_image, specified_tags, tagging_model_type, tagging_model, transform):
    print(f"Start processing, image size {raw_image.size}")

    image = transform(raw_image).unsqueeze(0).to(device)

    if tagging_model_type == "RAM":
        res = inference_ram(image, tagging_model)
        tags = res[0].strip(' ').replace('  ', ' ')
        tags_chinese = res[1].strip(' ').replace('  ', ' ')
        print("Tags: ", tags)
        print("标签: ", tags_chinese)
        return tags, tags_chinese
    else:
        res = inference_tag2text(image, tagging_model, specified_tags)
        tags = res[0].strip(' ').replace('  ', ' ')
        caption = res[2]
        print(f"Tags: {tags}")
        print(f"Caption: {caption}")
        return tags, caption


def inference_with_ram(img):
    return inference(img, None, "RAM", ram_model, transform)


def inference_with_t2t(img, input_tags):
    return inference(img, input_tags, "Tag2Text", tag2text_model, transform)


if __name__ == "__main__":
    import gradio as gr

    # get transform and load models
    transform = get_transform(image_size=image_size)
    ram_model = ram(pretrained=ram_checkpoint, image_size=image_size, vit='swin_l').eval().to(device)
    tag2text_model = tag2text(
        pretrained=tag2text_checkpoint, image_size=image_size, vit='swin_b').eval().to(device)

    # build GUI
    def build_gui():

        description = """
            <center><strong><font size='10'>Recognize Anything Model</font></strong></center>
            <br>
            <p>Welcome to the <a href='https://recognize-anything.github.io/' target='_blank'>Recognize Anything Model</a> / <a href='https://tag2text.github.io/Tag2Text' target='_blank'>Tag2Text Model</a> demo!</p>
            <li>
                <b>Recognize Anything Model:</b> Upload your image to get the <b>English and Chinese tags</b>!
            </li>
            <li>
                <b>Tag2Text Model:</b> Upload your image to get the <b>tags and caption</b>! (Optional: Specify tags to get the corresponding caption.)
            </li>
            <p><b>More over:</b> Combine with <a href='https://github.com/IDEA-Research/Grounded-Segment-Anything' target='_blank'>Grounded-SAM</a>, you can get <b>boxes and masks</b>! Please run <a href='https://github.com/xinyu1205/recognize-anything/blob/main/gui_demo.ipynb' target='_blank'>this notebook</a> to try out!</p>
            <p>Great thanks to <a href='https://huggingface.co/majinyu' target='_blank'>Ma Jinyu</a>, the major contributor of this demo!</p>
        """  # noqa

        article = """
            <p style='text-align: center'>
                RAM and Tag2Text are trained on open-source datasets, and we are persisting in refining and iterating upon it.<br/>
                <a href='https://recognize-anything.github.io/' target='_blank'>Recognize Anything: A Strong Image Tagging Model</a>
                |
                <a href='https://tag2text.github.io/' target='_blank'>Tag2Text: Guiding Language-Image Model via Image Tagging</a>
            </p>
        """  # noqa

        with gr.Blocks(title="Recognize Anything Model") as demo:
            ###############
            # components
            ###############
            gr.HTML(description)

            with gr.Tab(label="Recognize Anything Model"):
                with gr.Row():
                    with gr.Column():
                        ram_in_img = gr.Image(type="pil")
                        with gr.Row():
                            ram_btn_run = gr.Button(value="Run")
                            try:
                                ram_btn_clear = gr.ClearButton()
                            except AttributeError:  # old gradio does not have ClearButton, not big problem
                                ram_btn_clear = None
                    with gr.Column():
                        ram_out_tag = gr.Textbox(label="Tags")
                        ram_out_biaoqian = gr.Textbox(label="标签")
                gr.Examples(
                    examples=[
                        ["images/demo1.jpg"],
                        ["images/demo2.jpg"],
                        ["images/demo4.jpg"],
                    ],
                    fn=inference_with_ram,
                    inputs=[ram_in_img],
                    outputs=[ram_out_tag, ram_out_biaoqian],
                    cache_examples=True
                )

            with gr.Tab(label="Tag2Text Model"):
                with gr.Row():
                    with gr.Column():
                        t2t_in_img = gr.Image(type="pil")
                        t2t_in_tag = gr.Textbox(label="User Specified Tags (Optional, separated by comma)")
                        with gr.Row():
                            t2t_btn_run = gr.Button(value="Run")
                            try:
                                t2t_btn_clear = gr.ClearButton()
                            except AttributeError:  # old gradio does not have ClearButton, not big problem
                                t2t_btn_clear = None
                    with gr.Column():
                        t2t_out_tag = gr.Textbox(label="Tags")
                        t2t_out_cap = gr.Textbox(label="Caption")
                gr.Examples(
                    examples=[
                        ["images/demo4.jpg", ""],
                        ["images/demo4.jpg", "power line"],
                        ["images/demo4.jpg", "track, train"],
                    ],
                    fn=inference_with_t2t,
                    inputs=[t2t_in_img, t2t_in_tag],
                    outputs=[t2t_out_tag, t2t_out_cap],
                    cache_examples=True
                )

            gr.HTML(article)

            ###############
            # events
            ###############
            # run inference
            ram_btn_run.click(
                fn=inference_with_ram,
                inputs=[ram_in_img],
                outputs=[ram_out_tag, ram_out_biaoqian]
            )
            t2t_btn_run.click(
                fn=inference_with_t2t,
                inputs=[t2t_in_img, t2t_in_tag],
                outputs=[t2t_out_tag, t2t_out_cap]
            )

            # clear
            if ram_btn_clear is not None:
                ram_btn_clear.add([ram_in_img, ram_out_tag, ram_out_biaoqian])
            if t2t_btn_clear is not None:
                t2t_btn_clear.add([t2t_in_img, t2t_in_tag, t2t_out_tag, t2t_out_cap])

        return demo

    build_gui().launch(enable_queue=True)