File size: 7,074 Bytes
da50cbf | 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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 | import base64
import requests
from io import BytesIO, StringIO
from typing import Union, Optional, Tuple, List
from PIL import Image, ImageOps
import os
def get_api_key(file_path):
# Read the API key from the first line of the file
with open(file_path, 'r') as file:
return file.readline().strip()
# Function to encode the image
def encode_image(image_path):
with open(image_path, "rb") as image_file:
return base64.b64encode(image_file.read()).decode('utf-8')
def pick_next_item(current_item, item_list):
if current_item not in item_list:
raise ValueError("Current item is not in the list")
current_index = item_list.index(current_item)
next_index = (current_index + 1) % len(item_list)
return item_list[next_index]
# Function to encode a PIL image
def encode_pil_image(pil_image):
# Create an in-memory binary stream
image_stream = BytesIO()
# Save the PIL image to the binary stream in JPEG format (you can change the format if needed)
pil_image.save(image_stream, format='JPEG')
# Get the binary data from the stream and encode it as base64
image_data = image_stream.getvalue()
base64_image = base64.b64encode(image_data).decode('utf-8')
return base64_image
def load_image(image: Union[str, Image.Image], format: str = "RGB", size: Optional[Tuple] = None) -> Image.Image:
"""
Load an image from a given path or URL and convert it to a PIL Image.
Args:
image (Union[str, Image.Image]): The image path, URL, or a PIL Image object to be loaded.
format (str, optional): Desired color format of the resulting image. Defaults to "RGB".
size (Optional[Tuple], optional): Desired size for resizing the image. Defaults to None.
Returns:
Image.Image: A PIL Image in the specified format and size.
Raises:
ValueError: If the provided image format is not recognized.
"""
if isinstance(image, str):
if image.startswith("http://") or image.startswith("https://"):
image = Image.open(requests.get(image, stream=True).raw)
elif os.path.isfile(image):
image = Image.open(image)
else:
raise ValueError(
f"Incorrect path or url, URLs must start with `http://` or `https://`, and {image} is not a valid path"
)
elif isinstance(image, Image.Image):
image = image
else:
raise ValueError(
"Incorrect format used for image. Should be an url linking to an image, a local path, or a PIL image."
)
image = ImageOps.exif_transpose(image)
image = image.convert(format)
if (size != None):
image = image.resize(size, Image.LANCZOS)
return image
class GPT4v():
def __init__(self, key, url="https://api.openai.com/v1/chat/completions", are_images_encoded=False, model_name="gpt-4-vision-preview"):
"""OpenAI GPT-4-vision model wrapper
Args:
api_key_path (str): Path to the API key file. Defaults to 'keys/secret.env'.
are_images_encoded (bool): Whether the images are encoded in base64. Defaults to False.
"""
self.multiple_api_keys = False
self.current_key_file = None
self.api_key = key
self.url = url
self.model_name = model_name
self.use_encode = are_images_encoded
def prepare_input(self, image_links: List = [], text_prompt: str = ""):
prompt_content = []
text_dict = {
"type": "text",
"text": text_prompt
}
prompt_content.append(text_dict)
if not isinstance(image_links, list):
image_links = [image_links]
for image_link in image_links:
image = load_image(image_link)
if self.use_encode:
visual_dict = {
"type": "image_url",
"image_url": {"url": f"data:image/jpeg;base64,{encode_pil_image(image)}"}
}
else:
visual_dict = {
"type": "image_url",
"image_url": {"url": image_link}
}
prompt_content.append(visual_dict)
return prompt_content
def inference(self, prompt, seed: Optional[int] = None):
payload = {
"model": self.model_name,
"messages": [
{
"role": "user",
"content": prompt
}
],
# "max_tokens": 1400
}
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {self.api_key}"
}
# try:
response = requests.post(self.url, json=payload, headers=headers, timeout=180) # Set timeout to 5 minutes (300 seconds)
# except Exception as e:
# print(f"Error: {e}")
# return ""
#return response.text
return self.extract_response(response)
def extract_response(self, response):
try:
response = response.json()
out = response['choices'][0]['message']['content']
return out
except Exception as e:
print(f"Error: {e}")
if response['error']['code'] == 'content_policy_violation':
print("Code is content_policy_violation")
elif response['error']['code'] in ['rate_limit_exceeded', 'insufficient_quota', 'insufficient_user_quota']:
print(f"Code is {response['error']['code']}", flush=True)
print(response['error']['message'], flush=True)
return "rate_limit_exceeded"
if self.multiple_api_keys == True:
new_key = pick_next_item(self.current_key_file, self.key_lists)
self.update_key(new_key)
self.current_key_file = new_key #override key
print("New key is from the file: ", new_key)
else:
print("Code is different")
print(response)
print(f"{response['error']['code']=}")
return ""
def update_key(self, key, load_from_file=True):
if load_from_file:
self.api_key = get_api_key(key)
else:
self.api_key = key
class GPT4o(GPT4v):
def __init__(self, key, url="https://api.openai.com/v1/chat/completions", are_images_encoded=False, model_name="gpt-4o-2024-05-13"):
super().__init__(key, url, are_images_encoded, model_name)
if __name__ == "__main__":
model = GPT4o('sk-cB6h7HcCSDIp71gs6lFLZxKE0dOYOnJbxzES6kWXe1Wb2VHS', model_name="gpt-4.1")
prompt = model.prepare_prompt(['https://chromaica.github.io/Museum/ImagenHub_Text-Guided_IE/DiffEdit/sample_34_1.jpg', 'https://chromaica.github.io/Museum/ImagenHub_Text-Guided_IE/input/sample_34_1.jpg'], 'What is difference between two images?')
print("prompt : \n", prompt)
res = model.get_parsed_output(prompt)
print("result : \n", res) |