saljoofri's picture
Update app.py
e57b1cd
raw
history blame contribute delete
No virus
1.62 kB
##### Set API variables #####
import os
hf_api_key = os.environ['HF_API_KEY']
hf_api_itt_base = os.environ['HF_API_ITT_BASE']
##### Helper functions #####
import io, base64, requests, json
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 get_completion(inputs, parameters = None, ENDPOINT_URL = hf_api_itt_base):
headers = {
"Authorization": f"Bearer {hf_api_key}",
"Content-Type": "application/json"
}
data = { "inputs" : inputs }
if parameters is not None:
data.update({"parameters" : parameters})
response = requests.post(ENDPOINT_URL,
headers = headers,
data = json.dumps(data))
return json.loads(response.content.decode("utf-8"))
def captioner(pil_image):
base64_image = image_to_base64_str(pil_image)
result = get_completion(base64_image)
return result[0]['generated_text']
##### Gradio App #####
import gradio as gr
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 = ["tank-aerial-view.jpg", "f35-carrier-takeoff.jpg", "the_los_angelesclass_submarine.jpg"])
demo.launch()