tiktok caption generator from text to image

caption-generator.py

from PIL import Image, ImageDraw, ImageFont, ImageColor


class TextCaption:
    def __init__(self, **kwargs):
        self.kwargs = kwargs
        # Create a new image with a white background
        self.max_width = self.kwargs['max_width']

        # Load a font
        self.font_size = 60 if not 'font_size' in self.kwargs else self.kwargs['font_size']
        self.font = ImageFont.truetype(self.kwargs['font_file'], self.font_size)  # Use a specific font

        # spacing
        self.padding = 15 if not 'padding' in self.kwargs else self.kwargs['padding']
        self.line_spacing = 25 if not 'line_spacing' in self.kwargs else self.kwargs['line_spacing']

    def generate(self, file, text):
        wrapped_lines = self.wrap_text(text, self.font, self.max_width)

        list_width = []
        total_height = 0
        for line in wrapped_lines:
            width, height = self.get_size(self.font.getbbox(line))
            list_width.append(width + self.padding + (self.line_spacing * 2))
            total_height += height + self.padding + (self.line_spacing * 2)
        max_line_width = max(list_width)
        self.width = max_line_width
        self.height = total_height

        self.image = Image.new('RGBA', (self.width, self.height), (0, 0, 0, 0))  # Transparent background
        # Create a draw object
        self.draw = ImageDraw.Draw(self.image)

        # Define the text to be wrapped
        font_color = self.hex_to_rgb(self.kwargs['font_color'])  # Black text
        background_color = self.hex_to_rgb(self.kwargs['background_color'])  # White background for the text (optional)

        # Calculate the total height of the text block
        total_text_height = sum(
            [self.draw.textbbox((0, 0), line, font=self.font)[3] - self.draw.textbbox((0, 0), line, font=self.font)[1] for line in
             wrapped_lines]) + (len(wrapped_lines) - 1) * self.line_spacing + self.padding * 2

        # Calculate the starting vertical position to center the text block
        text_y = self.line_spacing + self.padding

        # Draw each line of text with its own rounded rectangle
        for line in wrapped_lines:
            # Calculate text size using textbbox
            bbox = self.draw.textbbox((0, 0), line, font=self.font)
            text_width = bbox[2] - bbox[0]
            text_height = self.font_size + self.line_spacing

            # Position the text
            text_x = (self.width - text_width) / 2

            # Draw a rounded rectangle behind the text
            rounded_rectangle_xy = (
                text_x - self.padding,
                text_y - self.padding + 5,
                text_x + text_width + self.padding,
                text_y + text_height + self.padding + 5
            )
            self.draw_rounded_rectangle(self.draw, rounded_rectangle_xy, radius=15, fill=background_color)

            # Draw the text on the image
            self.draw.text((text_x, text_y), line, font=self.font, fill=font_color)

            # Update the vertical position for the next line
            text_y += text_height + self.line_spacing

        # Save the image
        self.image.save(file, format='PNG')

    def show(self):
        # Optionally, show the image
        self.image.show()

    def hex_to_rgb(self, hex_color):
        # Use ImageColor to convert hex to RGB
        rgb_color = ImageColor.getcolor(hex_color, "RGB")
        return rgb_color

    def draw_rounded_rectangle(self, draw, xy, radius, fill):
        """Draw a rounded rectangle."""
        x0, y0, x1, y1 = xy
        draw.rectangle([x0 + radius, y0, x1 - radius, y1], fill=fill)  # Middle rectangle
        draw.rectangle([x0, y0 + radius, x1, y1 - radius], fill=fill)  # Middle rectangle
        draw.ellipse([x0, y0, x0 + radius * 2, y0 + radius * 2], fill=fill)  # Top-left corner
        draw.ellipse([x1 - radius * 2, y0, x1, y0 + radius * 2], fill=fill)  # Top-right corner
        draw.ellipse([x0, y1 - radius * 2, x0 + radius * 2, y1], fill=fill)  # Bottom-left corner
        draw.ellipse([x1 - radius * 2, y1 - radius * 2, x1, y1], fill=fill)  # Bottom-right corner

    def get_size(self, getbbox):
        left, top, right, bottom = getbbox
        text_width = right - left
        text_height = bottom - top
        return (text_width, text_height)

    def wrap_text(self, text, font, max_width):
        """Wrap text to fit within a specified width."""
        lines = []
        words = text.split(' ')
        current_line = ""

        for word in words:
            # Check the width of the current line plus the new word
            test_line = f"{current_line} {word}".strip()
            if self.get_size(font.getbbox(test_line))[0] <= max_width:
                current_line = test_line
            else:
                # If the line is too long, add the current line to the list and start a new line
                if current_line:
                    lines.append(current_line)
                current_line = word  # Start a new line with the current word

        # Add the last line if it exists
        if current_line:
            lines.append(current_line)

        return lines

if __name__ == '__main__':
    caption = TextCaption(
        max_width=540,
        font_color='black',
        background_color='white',
        font_file='test.ttf',
        font_size=35,
        line_spacing=10
    )
    caption.generate('output_image.png', 'Hal hal yang harus diperhatikan kalau fb mau monetisasi')
    caption.show()