feat: support many template in image generation (#13)

- feat: support many template in image generation (547fbf9dcedd839f06720760776af450e5633062)

.gitattributes

@@ -38,3 +38,5 @@ assets/Montserrat-Regular.ttf filter=lfs diff=lfs merge=lfs -text
 assets/template_natura_empty.jpg filter=lfs diff=lfs merge=lfs -text
 assets/template_1.png filter=lfs diff=lfs merge=lfs -text
 assets/template_2.png filter=lfs diff=lfs merge=lfs -text
+assets/template_b_avon.png filter=lfs diff=lfs merge=lfs -text
+assets/template_b_natura.png filter=lfs diff=lfs merge=lfs -text

app.py

@@ -132,6 +132,11 @@ with gr.Blocks() as demo:
                 image_original_price_input = gr.Number(label="Original Price", placeholder="Enter original price...")
                 image_final_price_input = gr.Number(label="Final Price", placeholder="Enter final price...")
                 image_coupon_code_input = gr.Textbox(label="Coupon Code", placeholder="Enter coupon code...")
+                template_selection = gr.Radio(
+                    choices=["lidi_promo", "natura", "avon"],
+                    value="lidi_promo",
+                    label="Template Selection"
+                )
                 gen_image_btn = gr.Button("Generate Image")
             with gr.Column():
                 image_output = gr.Image(label="Generated Image", height=500, type="filepath", interactive=False, show_share_button=True)
@@ -168,7 +173,7 @@ with gr.Blocks() as demo:
 
     
 
-    def generate_image(product_image_url, product_name, original_price, final_price, coupon_code):
+    def generate_image(product_image_url, product_name, original_price, final_price, coupon_code, template_name):
         tool = GenerateImageTool()
         original_price_str = f"{original_price:.2f}".replace('.', ',')
         final_price_str = f"{final_price:.2f}".replace('.', ',')
@@ -180,7 +185,8 @@ with gr.Blocks() as demo:
             product_name=product_name,
             original_price=original_price_str,
             final_price=final_price_str,
-            coupon_code=coupon_code
+            coupon_code=coupon_code,
+            template_name=template_name
         )
         
         yield gr.update(interactive=True, value="Generate Image"), image_path, gr.update(interactive=True)
@@ -207,7 +213,7 @@ with gr.Blocks() as demo:
                                      inputs=[fragrantica_url_input, openai_key_input, natura_token_input, openai_base_url_input, openai_model_name_input], 
                                      outputs=fragrantica_output)
     gen_image_btn.click(generate_image,
-                        inputs=[image_product_url_input, image_product_name_input, image_original_price_input, image_final_price_input, image_coupon_code_input],
+                        inputs=[image_product_url_input, image_product_name_input, image_original_price_input, image_final_price_input, image_coupon_code_input, template_selection],
                         outputs=[gen_image_btn, image_output, share_button])
     share_button.click(fn=process_image_for_sharing, 
                        inputs=[image_output], 

image_generator_tool.py

@@ -1,9 +1,6 @@
 from crewai.tools import BaseTool
 from pydantic import BaseModel, Field
-from PIL import Image, ImageDraw, ImageFont
-import requests
-from io import BytesIO
-import base64
+from template_system import TemplateRegistry
 
 class GenerateImageToolInput(BaseModel):
     """Input for the Generate Image Tool."""
@@ -12,63 +9,28 @@ class GenerateImageToolInput(BaseModel):
     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
+    template_name: str = Field(default="lidi_promo", description="Name of the template to use for generating the image.")
 
 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'
-
+    def _run(self, product_image_url: str, product_name: str, original_price: str, final_price: str, coupon_code: str, template_name: str = "lidi_promo") -> str:
         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."
+            # Get the template instance
+            template = TemplateRegistry.get_template(template_name)
+            
+            # Generate the image using the template
+            return template.generate_image(
+                product_image_url=product_image_url,
+                product_name=product_name,
+                original_price=original_price,
+                final_price=final_price,
+                coupon_code=coupon_code
+            )
+            
+        except ValueError as e:
+            return f"Error: {e}"
         except Exception as e:
             return f"An error occurred: {e}"

template_system.py

@@ -0,0 +1,359 @@
+from abc import ABC, abstractmethod
+from typing import Dict, Any
+from PIL import Image, ImageDraw, ImageFont
+import requests
+from io import BytesIO
+import tempfile
+
+
+class Template(ABC):
+    """
+    Abstract base class for image templates.
+    Each template defines its own configuration for box sizes, positions, colors, and fonts.
+    """
+    
+    def __init__(self, template_path: str = None):
+        self.template_path = template_path
+    
+    @abstractmethod
+    def get_box_config(self) -> Dict[str, Any]:
+        """Return box configuration including size and position for product image."""
+        pass
+    
+    @abstractmethod
+    def get_text_config(self) -> Dict[str, Dict[str, Any]]:
+        """Return text configuration including positions, colors, and fonts for all text elements."""
+        pass
+    
+    @abstractmethod
+    def get_font_config(self) -> Dict[str, Dict[str, Any]]:
+        """Return font configuration for different text elements."""
+        pass
+    
+    def load_template_image(self) -> Image.Image:
+        """Load and return the template image."""
+        return Image.open(self.template_path).convert("RGBA")
+    
+    def load_fonts(self) -> Dict[str, ImageFont.FreeTypeFont]:
+        """Load and return all required fonts."""
+        fonts = {}
+        font_config = self.get_font_config()
+        
+        for font_name, config in font_config.items():
+            try:
+                fonts[font_name] = ImageFont.truetype(config['path'], config['size'])
+            except IOError:
+                print(f"Font {config['path']} not found. Using default font.")
+                fonts[font_name] = ImageFont.load_default()
+        
+        return fonts
+    
+    def generate_image(self, product_image_url: str, product_name: str, 
+                      original_price: str, final_price: str, coupon_code: str) -> str:
+        """
+        Generate the promotional image using this template's configuration.
+        
+        Args:
+            product_image_url: URL of the product image
+            product_name: Name of the product
+            original_price: Original price of the product
+            final_price: Final price of the product
+            coupon_code: Coupon code to display
+            
+        Returns:
+            Path to the generated image file
+        """
+        try:
+            # Load template and fonts
+            template_image = self.load_template_image()
+            fonts = self.load_fonts()
+            
+            # Fetch and process product image
+            response = requests.get(product_image_url)
+            product_image_data = BytesIO(response.content)
+            product_image = Image.open(product_image_data).convert("RGBA")
+            
+            # Get box configuration
+            box_config = self.get_box_config()
+            box_size = box_config['size']
+            box_position = box_config['position']
+            
+            # Resize product image to fit within box while preserving aspect ratio
+            product_image_resized = product_image.copy()
+            product_image_resized.thumbnail(box_size)
+            
+            # Calculate position to center the image in the box
+            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)
+            
+            # Paste product image onto template
+            template_image.paste(product_image_resized, paste_position, product_image_resized)
+            
+            # Draw text elements
+            draw = ImageDraw.Draw(template_image)
+            text_config = self.get_text_config()
+            
+            # Draw each text element
+            for element_name, config in text_config.items():
+                text_content = self._get_text_content(element_name, product_name, 
+                                                    original_price, final_price, coupon_code)
+                position = config['position']
+                color = config['color']
+                font_name = config['font']
+                anchor = config.get('anchor', 'ms')
+                
+                draw.text(position, text_content, font=fonts[font_name], 
+                         fill=color, anchor=anchor)
+            
+            # Save the result
+            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 '{self.template_path}' was not found."
+        except Exception as e:
+            return f"An error occurred: {e}"
+    
+    def _get_text_content(self, element_name: str, product_name: str, 
+                         original_price: str, final_price: str, coupon_code: str) -> str:
+        """Get the actual text content for each text element."""
+        content_map = {
+            'product_name': product_name,
+            'original_price': f"De: R$ {original_price}",
+            'final_price': f"Por: R$ {final_price}",
+            'coupon_code': coupon_code
+        }
+        return content_map.get(element_name, '')
+
+
+class TemplateRegistry:
+    """Registry for managing different template types."""
+    
+    _templates = {}
+    
+    @classmethod
+    def register(cls, name: str, template_class):
+        """Register a template class."""
+        cls._templates[name] = template_class
+    
+    @classmethod
+    def get_template(cls, name: str) -> Template:
+        """Get a template instance by name."""
+        if name not in cls._templates:
+            raise ValueError(f"Template '{name}' not found")
+        template_instance = cls._templates[name]()
+        return template_instance
+    
+    @classmethod
+    def list_templates(cls) -> list:
+        """List all registered template names."""
+        return list(cls._templates.keys())
+
+
+class LidiPromoTemplate(Template):
+    """
+    Template implementation for Lidi promotional images.
+    Uses the original hardcoded values from the existing implementation.
+    """
+    
+    def __init__(self, template_path: str = None):
+        super().__init__(template_path or "assets/template_1.png")
+    
+    def get_box_config(self) -> Dict[str, Any]:
+        """Return box configuration for product image."""
+        return {
+            'size': (442, 353),
+            'position': (140, 280)  # (x, y) from top-left corner
+        }
+    
+    def get_text_config(self) -> Dict[str, Dict[str, Any]]:
+        """Return text configuration for all text elements."""
+        return {
+            'product_name': {
+                'position': (360, 710),
+                'color': '#FFFFFF',
+                'font': 'font_name',
+                'anchor': 'ms'
+            },
+            'original_price': {
+                'position': (360, 800),
+                'color': '#FFFFFF',
+                'font': 'font_price_from',
+                'anchor': 'ms'
+            },
+            'final_price': {
+                'position': (360, 860),
+                'color': '#FEE161',  # Yellow color from original design
+                'font': 'font_price',
+                'anchor': 'ms'
+            },
+            'coupon_code': {
+                'position': (360, 993),
+                'color': '#000000',
+                'font': 'font_cupom',
+                'anchor': 'ms'
+            }
+        }
+    
+    def get_font_config(self) -> Dict[str, Dict[str, Any]]:
+        """Return font configuration for different text elements."""
+        return {
+            'font_name': {
+                'path': 'assets/Montserrat-Bold.ttf',
+                'size': 47
+            },
+            'font_price_from': {
+                'path': 'assets/Montserrat-Regular.ttf',
+                'size': 28
+            },
+            'font_price': {
+                'path': 'assets/Montserrat-Bold.ttf',
+                'size': 47
+            },
+            'font_cupom': {
+                'path': 'assets/Montserrat-Bold.ttf',
+                'size': 33
+            }
+        }
+
+
+class NaturaTemplate(Template):
+    """
+    Template implementation for Natura promotional images.
+    Uses template_b_natura.png with different configuration.
+    """
+    
+    def __init__(self, template_path: str = None):   
+        super().__init__(template_path or "assets/template_b_natura.png")
+    
+    def get_box_config(self) -> Dict[str, Any]:
+        """Return box configuration for product image."""
+        return {
+            'size': (602, 424),
+            'position': (54, 254)  # (x, y) from top-left corner
+        }
+    
+    def get_text_config(self) -> Dict[str, Dict[str, Any]]:
+        """Return text configuration for all text elements."""
+        return {
+            'product_name': {
+                'position': (72, 727),
+                'color': '#000000',
+                'font': 'font_name',
+                'anchor': 'ls'
+            },
+            'original_price': {
+                'position': (72, 765),
+                'color': '#666666',
+                'font': 'font_price_from',
+                'anchor': 'ls'
+            },
+            'final_price': {
+                'position': (90, 837),
+                'color': '#FFFFFF',
+                'font': 'font_price',
+                'anchor': 'lm'
+            },
+            'coupon_code': {
+                'position': (461, 957),
+                'color': '#FFFFFF',
+                'font': 'font_cupom',
+                'anchor': 'ms'
+            }
+        }
+    
+    def get_font_config(self) -> Dict[str, Dict[str, Any]]:
+        """Return font configuration for different text elements."""
+        return {
+            'font_name': {
+                'path': 'assets/Montserrat-Bold.ttf',
+                'size': 32
+            },
+            'font_price_from': {
+                'path': 'assets/Montserrat-Regular.ttf',
+                'size': 22
+            },
+            'font_price': {
+                'path': 'assets/Montserrat-Bold.ttf',
+                'size': 40
+            },
+            'font_cupom': {
+                'path': 'assets/Montserrat-Bold.ttf',
+                'size': 42
+            }
+        }
+
+
+class AvonTemplate(Template):
+    """
+    Template implementation for Avon promotional images.
+    Uses template_b_avon.png with Avon-specific configuration.
+    """
+    
+    def __init__(self, template_path: str = None):
+        super().__init__(template_path or "assets/template_b_avon.png")
+    
+    def get_box_config(self) -> Dict[str, Any]:
+        """Return box configuration for product image."""
+        return {
+            'size': (602, 424),
+            'position': (54, 254)  # (x, y) from top-left corner
+        }
+    
+    def get_text_config(self) -> Dict[str, Dict[str, Any]]:
+        """Return text configuration for all text elements."""
+        return {
+            'product_name': {
+                'position': (72, 727),
+                'color': '#000000',
+                'font': 'font_name',
+                'anchor': 'ls'
+            },
+            'original_price': {
+                'position': (72, 765),
+                'color': '#666666',
+                'font': 'font_price_from',
+                'anchor': 'ls'
+            },
+            'final_price': {
+                'position': (90, 837),
+                'color': '#FFFFFF',
+                'font': 'font_price',
+                'anchor': 'lm'
+            },
+            'coupon_code': {
+                'position': (461, 957),
+                'color': '#FFFFFF',
+                'font': 'font_cupom',
+                'anchor': 'ms'
+            }
+        }
+    
+    def get_font_config(self) -> Dict[str, Dict[str, Any]]:
+        """Return font configuration for different text elements."""
+        return {
+            'font_name': {
+                'path': 'assets/Montserrat-Bold.ttf',
+                'size': 32
+            },
+            'font_price_from': {
+                'path': 'assets/Montserrat-Regular.ttf',
+                'size': 22
+            },
+            'font_price': {
+                'path': 'assets/Montserrat-Bold.ttf',
+                'size': 40
+            },
+            'font_cupom': {
+                'path': 'assets/Montserrat-Bold.ttf',
+                'size': 42
+            }
+        }
+
+
+# Register additional templates
+TemplateRegistry.register('lidi_promo', LidiPromoTemplate)
+TemplateRegistry.register('natura', NaturaTemplate)
+TemplateRegistry.register('avon', AvonTemplate)