File size: 4,351 Bytes
01602db | 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 | from __future__ import annotations
import os
import gradio as gr
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from gradio_client import Client, handle_file
DESCRIPTION = "# Comparing image captioning models"
ORIGINAL_SPACE_INFO = """\
- [GIT-large fine-tuned on COCO](https://huggingface.co/spaces/library-samples/image-captioning-with-git)
- [BLIP-large](https://huggingface.co/spaces/library-samples/image-captioning-with-blip)
- [BLIP-2 T5-XXL](https://huggingface.co/spaces/hysts/BLIP2-with-transformers)
"""
urllib.request.urlretrieve("http://images.cocodataset.org/val2017/000000039769.jpg", "cats.jpg")
urllib.request.urlretrieve(
"https://huggingface.co/datasets/nielsr/textcaps-sample/resolve/main/stop_sign.png", "stop_sign.png"
)
urllib.request.urlretrieve(
"https://cdn.openai.com/dall-e-2/demos/text2im/astronaut/horse/photo/0.jpg", "astronaut.jpg"
)
def generate_caption_git(image_path: str, headers: dict | None = None) -> str:
try:
client = Client("library-samples/image-captioning-with-git", headers=headers, verbose=False)
return client.predict(handle_file(image_path), api_name="/caption")
except Exception as e:
print(f"caption error: {e!r}")
gr.Warning("The GIT-large Space is currently unavailable. Please try again later.")
return ""
def generate_caption_blip(image_path: str, headers: dict | None = None) -> str:
try:
client = Client("library-samples/image-captioning-with-blip", headers=headers, verbose=False)
return client.predict(handle_file(image_path), "A picture of", api_name="/caption")
except Exception as e:
print(f"caption error: {e!r}")
gr.Warning("The BLIP-large Space is currently unavailable. Please try again later.")
return ""
def generate_caption_blip2_t5xxl(image_path: str, headers: dict | None = None) -> str:
try:
client = Client("hysts/BLIP2-with-transformers", headers=headers, verbose=False)
return client.predict(
handle_file(image_path),
"Beam search",
1, # temperature
1, # length penalty
1.5, # repetition penalty
50, # max length
1, # min length
5, # number of beams
0.9, # top p
api_name="/caption",
)
except Exception as e:
print(f"caption error: {e!r}")
gr.Warning("The BLIP2 T5-XXL Space is currently unavailable. Please try again later.")
return ""
def generate_captions(image_path: str, request: gr.Request | None = None) -> tuple[str, str, str]:
# Forward the visitor's ZeroGPU identity so the upstream Spaces bill their quota, not this Space's IP.
headers = None
if request is not None and request.headers.get("x-ip-token"):
headers = {"X-IP-Token": request.headers["x-ip-token"]}
fns = (generate_caption_git, generate_caption_blip, generate_caption_blip2_t5xxl)
with ThreadPoolExecutor(max_workers=len(fns)) as pool:
return tuple(pool.map(lambda fn: fn(image_path, headers), fns))
with gr.Blocks() as demo:
gr.Markdown(DESCRIPTION)
with gr.Row():
with gr.Column():
input_image = gr.Image(type="filepath")
run_button = gr.Button("Caption")
with gr.Column():
out_git = gr.Textbox(label="GIT-large fine-tuned on COCO")
out_blip = gr.Textbox(label="BLIP-large")
out_blip2_t5xxl = gr.Textbox(label="BLIP-2 T5-XXL")
outputs = [
out_git,
out_blip,
out_blip2_t5xxl,
]
gr.Examples(
examples=[
"cats.jpg",
"stop_sign.png",
"astronaut.jpg",
],
inputs=input_image,
outputs=outputs,
fn=generate_captions,
cache_examples=os.getenv("CACHE_EXAMPLES") == "1",
)
with gr.Accordion(label="The original Spaces can be found here:", open=False):
gr.Markdown(ORIGINAL_SPACE_INFO)
run_button.click(
fn=generate_captions,
inputs=input_image,
outputs=outputs,
api_name="caption",
)
if __name__ == "__main__":
demo.queue(max_size=20).launch(css_paths="style.css")
|