Spaces:
Sleeping
Sleeping
| from transformers import pipeline, AutoTokenizer | |
| import io | |
| import base64 | |
| from PIL import Image | |
| import gradio as gr | |
| model = "Salesforce/blip-image-captioning-large" | |
| tokenizer = AutoTokenizer.from_pretrained(model, use_fast=True) | |
| pipe = pipeline(task="image-to-text", | |
| model=model, | |
| tokenizer=tokenizer) | |
| def image_to_base64(image: Image) -> str: | |
| """ | |
| Convert an image to a base64 string. | |
| """ | |
| bytearray= io.BytesIO() | |
| image.save(bytearray, format="PNG") | |
| return str(base64.b64encode(bytearray.getvalue()).decode('utf-8')) | |
| def caption_image(image): | |
| result = pipe( | |
| image_to_base64(image), | |
| #Temperature=0.7, | |
| # max_length=130, | |
| # min_length=30, | |
| #do_sample=True | |
| ) | |
| return result[0]['generated_text'].upper() | |
| if __name__ == "__main__": | |
| gr.close_all() | |
| with gr.Blocks() as interface: | |
| gr.Markdown("### Image Captioning using BLIP Large") | |
| with gr.Row(): | |
| image_input = gr.Image(type="pil", label="Image") | |
| with gr.Row(): | |
| caption_output = gr.Textbox(lines=2, label="Caption") | |
| with gr.Row(): | |
| clear_button = gr.ClearButton() | |
| caption_button = gr.Button("Caption", variant="primary") | |
| with gr.Row(): | |
| example_images = gr.Examples( | |
| examples=[ | |
| "data/image1.jpg", | |
| "data/image2.png", | |
| "data/image3.jpg", | |
| "data/image4.jpg", | |
| "data/image5.jpg", | |
| "data/image6.png", | |
| "data/image7.png", | |
| "data/image8.jpeg", | |
| "data/image9.jpeg", | |
| "data/image10.jpg", | |
| ], | |
| inputs=[image_input], | |
| label="Example Images" | |
| ) | |
| caption_button.click(fn=caption_image, | |
| inputs=[image_input], | |
| outputs=[caption_output] | |
| ) | |
| clear_button.click(fn=lambda: [None,""], | |
| inputs=[], | |
| outputs=[image_input, caption_output]) | |
| interface.launch(share=True, server_port=7860) | |