Spaces:
Runtime error
Runtime error
File size: 6,841 Bytes
8bdb03f 9027584 8bdb03f 9027584 43b47b1 8bdb03f b683f14 9027584 8bdb03f 9027584 8bdb03f 9027584 8bdb03f 9027584 8bdb03f 9027584 8bdb03f 9027584 8bdb03f 9027584 8bdb03f 9027584 8bdb03f 9027584 8bdb03f 9027584 8bdb03f 9027584 2c0c15f 9027584 8bdb03f 2c0c15f 8bdb03f 9027584 8bdb03f 9027584 8bdb03f 9027584 8bdb03f 9027584 8bdb03f 9027584 8bdb03f 9027584 8bdb03f 9027584 8bdb03f 9027584 8bdb03f 9027584 2c0c15f 8bdb03f 9027584 0b46e91 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 |
import torch
from ram import get_transform, inference_ram, inference_tag2text
from ram.models import ram, tag2text_caption
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_caption(
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://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)
|