File size: 1,643 Bytes
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
#!/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

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()
    img_obj.save(buffered, format="JPEG")
    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
    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))