nat-ad / image_generator_tool.py
ibombonato's picture
Upload image_generator_tool.py (#6)
f599b5b verified
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}"