nat-ad / image_generator_tool.py
ibombonato's picture
Add Mercado Livre support
472175a verified
raw
history blame
3.5 kB
from crewai.tools import BaseTool
from pydantic import BaseModel, Field
from PIL import Image, ImageDraw, ImageFont
import requests
from io import BytesIO
import uuid
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
import os
class GenerateImageTool(BaseTool):
name: str = "Generate Image Tool"
description: str = "Generates a promotional image for a product using a template."
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_natura_empty.jpg'
temp_dir = tempfile.gettempdir()
output_path = os.path.join(temp_dir, f'{uuid.uuid4()}.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), original_price, font=font_price_from, fill=white_color, anchor="ms")
draw.text((360, 860), final_price, font=font_price, fill=yellow_color, anchor="ms")
draw.text((360, 993), coupon_code, font=font_cupom, fill=black_color, anchor="ms")
template_image.save(output_path)
return output_path
except FileNotFoundError:
return f"Error: The template file '{template_path}' was not found."
except Exception as e:
return f"An error occurred: {e}"