Spaces:
Sleeping
Sleeping
File size: 3,169 Bytes
520da1b | 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 | # 1. 扫描当前目录下所有image文件 包含png jpg jpeg webp，进行编号 存储下list
import os, json
image_pathes = []
root = './images/'
image_path_file = 'image_pathes.json'
if os.path.exists(image_path_file):
with open(image_path_file, 'r') as f:
image_pathes = json.load(f)
print(len(image_pathes))
else:
print('error')
from openai import OpenAI
from PIL import Image
import base64
from io import BytesIO
client = OpenAI(
api_key=os.getenv("OPENAI_API_KEY") or os.getenv("AIHUBMIX_API_KEY", ""),
base_url=os.getenv("OPENAI_BASE_URL", "https://aihubmix.com/v1")
)
def resize_image(img, max_size=512):
width, height = img.size
ratio = min(max_size / width, max_size / height)
if ratio >= 1:
return img
new_width = int(width * ratio)
new_height = int(height * ratio)
resized_img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
return resized_img
def repaint_image(img):
# rapaint transparent area with white color
img = img.convert('RGBA')
data = img.getdata()
new_data = []
for item in data:
if item[3] == 0:
new_data.append((255, 255, 255, 255))
else:
new_data.append(item)
img.putdata(new_data)
img = img.convert('RGB')
# img.save('temp.png')
return img
def image_to_base64(image_path, show=False, target_size=512):
with Image.open(image_path) as img:
# img = img.resize(size)
img = resize_image(img, 512)
img = repaint_image(img)
buffered = BytesIO()
img.save(buffered, format="PNG")
img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
return img_base64
wwxxhh = 0
def ask_image(prompt, image_data):
global wwxxhh
number_of_trials = 0
while number_of_trials < 5:
try:
response = client.chat.completions.create(
# model="gpt-4o-mini",
model="gemini-2.0-flash",
messages=[
{
"role": "user",
"content": [
{
"type": "text",
"text": prompt},
{
"type": "image_url",
"image_url": {
"url": f"data:image/jpeg;base64,{image_data}"
},
},
],
}
]
)
wwxxhh += response.usage.total_tokens
return response.choices[0].message.content
except Exception as e:
number_of_trials += 1
print(e)
return 'Error!'
# 2. 读取prompt.json文件，读取整个作为字符串，逐个读取image文件，调用ask_image函数，将返回的结果存储下来
import json
with open('prompt.json', 'r') as f:
prompt = f.read()
# print(prompt)
ques_id = 15394
results_path = './results/'
image_path = image_pathes[ques_id]
print(image_path)
image_data = image_to_base64(image_path)
result = ask_image(prompt, image_data)
print(result)
# from IPython import embed
# embed()
|