File size: 3,991 Bytes
52b089b |
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 |
import base64
import json
import mimetypes
import os
import uuid
from io import BytesIO
from typing import Optional
import requests
from dotenv import load_dotenv
from PIL import Image
from smolagents import Tool, tool
load_dotenv(override=True)
def encode_image(image_path):
if image_path.startswith("http"):
user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36 Edg/119.0.0.0"
request_kwargs = {
"headers": {"User-Agent": user_agent},
"stream": True,
}
# Send a HTTP request to the URL
response = requests.get(image_path, **request_kwargs)
response.raise_for_status()
content_type = response.headers.get("content-type", "")
extension = mimetypes.guess_extension(content_type)
if extension is None:
extension = ".download"
fname = str(uuid.uuid4()) + extension
download_path = os.path.abspath(os.path.join("downloads", fname))
with open(download_path, "wb") as fh:
for chunk in response.iter_content(chunk_size=512):
fh.write(chunk)
image_path = download_path
with open(image_path, "rb") as image_file:
return base64.b64encode(image_file.read()).decode("utf-8")
def resize_image(image_path):
img = Image.open(image_path)
width, height = img.size
img = img.resize((int(width / 2), int(height / 2)))
new_image_path = f"resized_{image_path}"
img.save(new_image_path)
return new_image_path
@tool
def visualizer(image_path: str, question: Optional[str] = None) -> str:
"""A tool that can answer questions about attached images.
Args:
image_path: The path to the image on which to answer the question. This should be a local path to downloaded image.
question: The question to answer.
"""
if not isinstance(image_path, str):
raise Exception("You should provide at least `image_path` string argument to this tool!")
add_note = False
if not question:
add_note = True
question = "Please write a detailed caption for this image."
mime_type, _ = mimetypes.guess_type(image_path)
base64_image = encode_image(image_path)
# Configuración para Ollama
model_id = os.getenv("MODEL_ID", "qwen2.5-coder:3b")
api_base = os.getenv("OPENAI_API_BASE", "http://localhost:11434/v1")
api_key = os.getenv("OPENAI_API_KEY", "ollama")
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {api_key}"
}
payload = {
"model": model_id,
"messages": [
{
"role": "user",
"content": [
{"type": "text", "text": question},
{"type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{base64_image}"}},
],
}
],
"max_tokens": 1000,
}
try:
response = requests.post(f"{api_base}/chat/completions", headers=headers, json=payload)
response.raise_for_status()
output = response.json()["choices"][0]["message"]["content"]
except Exception as e:
print(f"Error processing image: {str(e)}")
if "Payload Too Large" in str(e):
new_image_path = resize_image(image_path)
base64_image = encode_image(new_image_path)
payload["messages"][0]["content"][1]["image_url"]["url"] = f"data:{mime_type};base64,{base64_image}"
response = requests.post(f"{api_base}/chat/completions", headers=headers, json=payload)
response.raise_for_status()
output = response.json()["choices"][0]["message"]["content"]
else:
raise Exception(f"Error processing image: {str(e)}")
if add_note:
output = f"You did not provide a particular question, so here is a detailed caption for the image: {output}"
return output
|