lib/image_captioning.py

import io
import os

import cv2
import PIL
import requests

class ImageCaptioning:
    """
    Performing an API call to BLIP's huggingface inference API
    """
    def __init__(self):
        self.api_endpoint = os.environ["blip_api_url"]
        self.org_token = os.environ["auth_token"]
        self.headers = { "Authorization": f"Bearer {self.org_token}" }
        self.session = requests.session()
    
    def convert_to_bytes(self, image: PIL.Image.Image):
        data = io.BytesIO()
        image.save(data, format="PNG")
        return data.getvalue()

    def query(self, image: PIL.Image.Image):
        response = self.session.post(
            self.api_endpoint,
            headers=self.headers,
            data=self.convert_to_bytes(image)
        )
        return response.json()