nat-ad / image_generator_tool.py
ibombonato's picture
Upload image_generator_tool.py (#6)
f599b5b verified
raw
history blame
3.51 kB
from crewai.tools import BaseTool
from pydantic import BaseModel, Field
from PIL import Image, ImageDraw, ImageFont
import requests
from io import BytesIO
import base64
class GenerateImageToolInput(BaseModel):
"""Input for the Generate Image Tool."""
product_image_url: str = Field(..., description="URL of the product image to be placed on the template.")
product_name: str = Field(..., description="Name of the product.")
original_price: str = Field(..., description="Original price of the product.")
final_price: str = Field(..., description="Final price of the product.")
coupon_code: str = Field(..., description="Coupon code to be displayed on the image.")
import tempfile
class GenerateImageTool(BaseTool):
name: str = "Generate Image Tool"
description: str = "Generates a promotional image for a product using a template and returns the file path."
args_schema = GenerateImageToolInput
def _run(self, product_image_url: str, product_name: str, original_price: str, final_price: str, coupon_code: str) -> str:
template_path = 'assets/template_1.png'
try:
template_image = Image.open(template_path).convert("RGBA")
response = requests.get(product_image_url)
product_image_data = BytesIO(response.content)
product_image = Image.open(product_image_data).convert("RGBA")
box_size = (442, 353)
box_position = (140, 280)
product_image_resized = product_image.copy()
product_image_resized.thumbnail(box_size)
paste_x = box_position[0] + (box_size[0] - product_image_resized.width) // 2
paste_y = box_position[1] + (box_size[1] - product_image_resized.height) // 2
paste_position = (paste_x, paste_y)
template_image.paste(product_image_resized, paste_position, product_image_resized)
draw = ImageDraw.Draw(template_image)
try:
font_name = ImageFont.truetype("assets/Montserrat-Bold.ttf", 47)
font_price_from = ImageFont.truetype("assets/Montserrat-Regular.ttf", 28)
font_price = ImageFont.truetype("assets/Montserrat-Bold.ttf", 47)
font_cupom = ImageFont.truetype("assets/Montserrat-Bold.ttf", 33)
except IOError:
print("Arial font not found. Using default font.")
font_name = ImageFont.load_default()
font_price_from = ImageFont.load_default()
font_price = ImageFont.load_default()
font_cupom = ImageFont.load_default()
white_color = "#FFFFFF"
yellow_color = "#FEE161"
black_color = "#000000"
draw.text((360, 710), product_name, font=font_name, fill=white_color, anchor="ms")
draw.text((360, 800), f"De: R$ {original_price}", font=font_price_from, fill=white_color, anchor="ms")
draw.text((360, 860), f"Por: R$ {final_price}", font=font_price, fill=yellow_color, anchor="ms")
draw.text((360, 993), coupon_code, font=font_cupom, fill=black_color, anchor="ms")
with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file:
template_image.save(temp_file.name)
return temp_file.name
except FileNotFoundError:
return f"Error: The template file '{template_path}' was not found."
except Exception as e:
return f"An error occurred: {e}"