lv_recipe_chatbot/openai_vision.py

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_openai_vision.ipynb.

# %% auto 0
__all__ = ['SAMPLE_IMG_DIR', 'prompt1', 'extract_veg_ingredients', 'encode_image', 'format_image', 'get_gptv_completion']

# %% ../nbs/04_openai_vision.ipynb 3
import base64
import json
import logging
import os
import statistics
import time
from functools import partial
from glob import glob
from logging.handlers import RotatingFileHandler
from pathlib import Path

import numpy as np
import requests
from PIL import Image

from openai import OpenAI
import constants

# %% ../nbs/04_openai_vision.ipynb 7
SAMPLE_IMG_DIR = Path(f"{constants.ROOT_DIR}/assets/images/vegan_ingredients")

# %% ../nbs/04_openai_vision.ipynb 8
def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")

# %% ../nbs/04_openai_vision.ipynb 9
# fmt: off
def format_image(
    image: str, # Image file path
    size: int = 300
):
    # fmt: on
    img = Image.open(image)
    width, height = img.size
    ratio = min(size / width, size / height)
    width_new, height_new = (round(width * ratio), round(height * ratio))
    width_new = int(np.round(width_new / 64.0)) * 64
    height_new = int(np.round(height_new / 64.0)) * 64
    img = img.resize((width_new, height_new))
    img = img.convert("RGB")
    return img

# %% ../nbs/04_openai_vision.ipynb 17
# https://platform.openai.com/docs/guides/vision/low-or-high-fidelity-image-understanding
def get_gptv_completion(prompt: str, img: str, detail="low"):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {os.environ['OPENAI_API_KEY']}",
    }

    payload = {
        "model": "gpt-4o",
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{encode_image(img)}",
                            "detail": detail,
                        },
                    },
                ],
            }
        ],
        "max_tokens": 300,
    }

    response = requests.post(
        "https://api.openai.com/v1/chat/completions", headers=headers, json=payload
    )
    return response.json()["choices"][0]["message"]["content"]

# %% ../nbs/04_openai_vision.ipynb 27
prompt1 = "What vegan ingredients are in the image? Ouput as JSON list."

# %% ../nbs/04_openai_vision.ipynb 34
extract_veg_ingredients = partial(get_gptv_completion, prompt1)