AjithBharadwaj's picture
Create app.py
b1f99b2
raw
history blame
1.07 kB
from transformers import pipeline
import gradio as gr
import os
import io
import IPython.display
from PIL import Image
import base64
import gradio as gr
get_completion = pipeline("image-to-text",model="Salesforce/blip-image-captioning-base")
def image_to_base64_str(pil_image):
byte_arr = io.BytesIO()
pil_image.save(byte_arr, format='PNG')
byte_arr = byte_arr.getvalue()
return str(base64.b64encode(byte_arr).decode('utf-8'))
def captioner(image):
base64_image = image_to_base64_str(image)
result = get_completion(base64_image)
return result[0]['generated_text']
gr.close_all()
demo = gr.Interface(fn=captioner,
inputs=[gr.Image(label="Upload image", type="pil")],
outputs=[gr.Textbox(label="Caption")],
title="Image Captioning with BLIP",
description="Caption any image using the BLIP model",
allow_flagging="never",
examples=["christmas_dog.jpeg", "bird_flight.jpeg", "cow.jpeg"])
demo.launch(infline=False)