Spaces:
Running
Running
| import numpy as np | |
| import random | |
| import torch | |
| import torchvision.transforms as transforms | |
| from PIL import Image | |
| from models.tag2text import tag2text_caption, ram | |
| import gradio as gr | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| image_size = 384 | |
| normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225 | |
| ]) | |
| transform = transforms.Compose([transforms.Resize((image_size, image_size)),transforms.ToTensor(),normalize]) | |
| #######Tag2Text Model | |
| pretrained = 'tag2text_swin_14m.pth' | |
| model_tag2text = tag2text_caption(pretrained=pretrained, image_size=image_size, vit='swin_b' ) | |
| model_tag2text.eval() | |
| model_tag2text = model_tag2text.to(device) | |
| #######RAM Model | |
| pretrained = 'ram_swin_large_14m.pth' | |
| model_ram = ram(pretrained=pretrained, image_size=image_size, vit='swin_l' ) | |
| model_ram.eval() | |
| model_ram = model_ram.to(device) | |
| def inference(raw_image, model_n , input_tag): | |
| raw_image = raw_image.resize((image_size, image_size)) | |
| image = transform(raw_image).unsqueeze(0).to(device) | |
| if model_n == 'Recognize Anything Model': | |
| model = model_ram | |
| with torch.no_grad(): | |
| tags, tags_chinese = model.generate_tag(image) | |
| return tags[0],tags_chinese[0], 'none' | |
| else: | |
| model = model_tag2text | |
| model.threshold = 0.68 | |
| if input_tag == '' or input_tag == 'none' or input_tag == 'None': | |
| input_tag_list = None | |
| else: | |
| input_tag_list = [] | |
| input_tag_list.append(input_tag.replace(',',' | ')) | |
| with torch.no_grad(): | |
| caption, tag_predict = model.generate(image,tag_input = input_tag_list,max_length = 50, return_tag_predict = True) | |
| if input_tag_list == None: | |
| tag_1 = tag_predict | |
| tag_2 = ['none'] | |
| else: | |
| _, tag_1 = model.generate(image,tag_input = None, max_length = 50, return_tag_predict = True) | |
| tag_2 = tag_predict | |
| return tag_1[0],'none',caption[0] | |
| def build_gui(): | |
| description = """ | |
| <center><strong><font size='10'>Recognize Anything Model</font></strong></center> | |
| <br> | |
| Welcome to the Recognize Anything Model (RAM) and Tag2Text Model demo! <br><br> | |
| <li> | |
| <b>Recognize Anything Model:</b> Upload your image to get the <b>English and Chinese outputs of the image tags</b>! | |
| </li> | |
| <li> | |
| <b>Tag2Text Model:</b> Upload your image to get the <b>tags</b> and <b>caption</b> of the image. | |
| Optional: You can also input specified tags to get the corresponding caption. | |
| </li> | |
| """ # noqa | |
| article = """ | |
| <p style='text-align: center'> | |
| RAM and Tag2Text is training 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://https://tag2text.github.io/' target='_blank'>Tag2Text: Guiding Language-Image Model via Image Tagging</a> | |
| | | |
| <a href='https://github.com/xinyu1205/Tag2Text' target='_blank'>Github Repo</a> | |
| </p> | |
| """ # noqa | |
| def inference_with_ram(img): | |
| res = inference(img, "Recognize Anything Model", None) | |
| return res[0], res[1] | |
| def inference_with_t2t(img, input_tags): | |
| res = inference(img, "Tag2Text Model", input_tags) | |
| return res[0], res[2] | |
| 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") | |
| ram_btn_clear = gr.Button(value="Clear") | |
| 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") | |
| t2t_btn_clear = gr.Button(value="Clear") | |
| 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] | |
| ) | |
| # # images of two image panels should keep the same | |
| # # and clear old outputs when image changes | |
| # # slow due to internet latency when deployed on huggingface, comment out | |
| # def sync_img(v): | |
| # return [gr.update(value=v)] + [gr.update(value="")] * 4 | |
| # ram_in_img.upload(fn=sync_img, inputs=[ram_in_img], outputs=[ | |
| # t2t_in_img, ram_out_tag, ram_out_biaoqian, t2t_out_tag, t2t_out_cap | |
| # ]) | |
| # ram_in_img.clear(fn=sync_img, inputs=[ram_in_img], outputs=[ | |
| # t2t_in_img, ram_out_tag, ram_out_biaoqian, t2t_out_tag, t2t_out_cap | |
| # ]) | |
| # t2t_in_img.clear(fn=sync_img, inputs=[t2t_in_img], outputs=[ | |
| # ram_in_img, ram_out_tag, ram_out_biaoqian, t2t_out_tag, t2t_out_cap | |
| # ]) | |
| # t2t_in_img.upload(fn=sync_img, inputs=[t2t_in_img], outputs=[ | |
| # ram_in_img, ram_out_tag, ram_out_biaoqian, t2t_out_tag, t2t_out_cap | |
| # ]) | |
| # clear all | |
| def clear_all(): | |
| return [gr.update(value=None)] * 2 + [gr.update(value="")] * 5 | |
| ram_btn_clear.click(fn=clear_all, inputs=[], outputs=[ | |
| ram_in_img, t2t_in_img, | |
| ram_out_tag, ram_out_biaoqian, t2t_in_tag, t2t_out_tag, t2t_out_cap | |
| ]) | |
| t2t_btn_clear.click(fn=clear_all, inputs=[], outputs=[ | |
| ram_in_img, t2t_in_img, | |
| ram_out_tag, ram_out_biaoqian, t2t_in_tag, t2t_out_tag, t2t_out_cap | |
| ]) | |
| return demo | |
| if __name__ == "__main__": | |
| demo = build_gui() | |
| demo.launch(enable_queue=True) | |