| 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,
|
| 1,
|
| 1.5,
|
| 50,
|
| 1,
|
| 5,
|
| 0.9,
|
| 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]:
|
|
|
| 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")
|
|
|
