Spaces:
Runtime error
Runtime error
File size: 2,171 Bytes
17a6d8e 906f9ca 17a6d8e 906f9ca 17a6d8e bc6e7dd b9083a9 17a6d8e |
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 |
#!/usr/bin/env python
"""get_caption.py: """
__author__ = "Rishabh Gupta"
__copyright__ = "Copyright 2023, Zuma"
__created_date__ = "04-08-2023"
import os
import json
import requests
import base64
from io import BytesIO
import PIL
def get_image_caption(img_obj, prompt=None):
"""
Args:
img_obj: Base 64 image object
prompt: Prompt to be used for captioning
Returns:
Captions for the image
"""
# print("Image Obj", img_obj)
buffered = BytesIO()
if img_obj is None:
return "No image data provided!"
try:
img_obj.save(buffered, format="JPEG")
except PIL.UnidentifiedImageError as e:
print("Error in saving the image", e)
return "Invalid image format!"
except Exception as e:
print("Some error occurred while saving the image", e)
return "Some error occurred while loading the image!"
img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
url = os.environ.get("IMAGE_CAPTION_ENDPOINT")
auth_token = os.environ.get("IMAGE_CAPTION_AUTH_TOKEN")
# print("img_str", img_str)
payload = json.dumps({
"inputs": img_str,
"prompt": None # Prompts are disabled as accuracy is not very good or not experimented enough.
})
headers = {
'Authorization': auth_token,
'Content-Type': 'application/json'
}
response = requests.request("POST", url, headers=headers, data=payload)
print("Response Status Code = ", response.status_code)
print("Response = ", response.text)
if response.status_code == 200:
caption = json.loads(response.text)["captions"]
print(f"Caption = {caption}")
return caption
elif response.status_code == 502:
print("Status 502")
return "Model is getting loaded, please wait for 3-4 minutes and then try again."
else:
return "Error in generating the caption!"
if __name__ == '__main__':
with open("/Users/rishabh/Downloads/5dbadb934042f1.16205085439.jpg", "rb") as image_file:
encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
print(get_image_caption(encoded_string))
|