import discord
from discord.ext import commands
from discord import app_commands, File
from PIL import Image, ImageDraw, ImageFont, ImageSequence
import requests
import io
import os
import textwrap # Import textwrap for text wrapping
class CaptionCog(commands.Cog, name="Caption"):
"""Cog for captioning GIFs"""
# Define constants for magic numbers
CAPTION_PADDING = 10
DEFAULT_GIF_DURATION = 100
MIN_FONT_SIZE = 10
MAX_FONT_SIZE = 30 # Decreased max font size
TEXT_COLOR = (0, 0, 0) # Black text
BAR_COLOR = (255, 255, 255) # White bar
def __init__(self, bot):
self.bot = bot
# Define preferred font names/paths
self.preferred_fonts = [
os.path.join("FONT", "OPTIFutura-ExtraBlackCond.otf") # Bundled fallback
]
def _add_text_to_gif(self, image_bytes: bytes, caption_text: str):
"""
Adds text to each frame of a GIF.
The text is placed in a white bar at the top of the GIF.
"""
try:
gif = Image.open(io.BytesIO(image_bytes))
frames = []
# Determine font size (e.g., 20% of image height, capped)
font_size = max(
self.MIN_FONT_SIZE, min(self.MAX_FONT_SIZE, int(gif.height * 0.2))
)
font = None
for font_choice in self.preferred_fonts:
try:
font = ImageFont.truetype(font_choice, font_size)
print(f"Successfully loaded font: {font_choice}")
break
except IOError:
print(f"Could not load font: {font_choice}. Trying next option.")
if font is None:
print(
"All preferred fonts failed to load. Using Pillow's default font."
)
font = ImageFont.load_default()
# Adjust font size for default font if necessary, as it might render differently.
# This might require re-calculating text_width and text_height if default font is used.
# Calculate max text width based on image width and padding
max_text_width = gif.width - (2 * self.CAPTION_PADDING)
# Wrap text based on max width
# Estimate characters per line based on font size and image width
# This is a heuristic and might need adjustment based on the font
estimated_char_width = font_size * 0.6
if (
estimated_char_width == 0
): # Avoid division by zero if font_size is somehow 0
estimated_char_width = 1
chars_per_line = int(max_text_width / estimated_char_width)
if chars_per_line <= 0: # Ensure at least one character per line
chars_per_line = 1
wrapped_text = textwrap.wrap(caption_text, width=chars_per_line)
# Calculate total text height and bar height
# Create a dummy draw object to measure text height per line
dummy_image = Image.new("RGB", (1, 1))
dummy_draw = ImageDraw.Draw(dummy_image)
line_heights = []
for line in wrapped_text:
if hasattr(dummy_draw, "textbbox"):
text_bbox = dummy_draw.textbbox((0, 0), line, font=font)
line_heights.append(text_bbox[3] - text_bbox[1])
else: # For older Pillow versions, use textsize (deprecated)
line_heights.append(dummy_draw.textsize(line, font=font)[1])
total_text_height = sum(line_heights)
bar_height = total_text_height + (2 * self.CAPTION_PADDING)
for frame in ImageSequence.Iterator(gif):
frame = frame.convert("RGBA")
# Create a new image for the frame with space for the text bar
new_frame_width = frame.width
new_frame_height = frame.height + bar_height
new_frame = Image.new(
"RGBA", (new_frame_width, new_frame_height), (0, 0, 0, 0)
) # Transparent background for the new area
# Draw the white bar
draw = ImageDraw.Draw(new_frame)
draw.rectangle(
[(0, 0), (new_frame_width, bar_height)], fill=self.BAR_COLOR
)
# Paste the original frame below the bar
new_frame.paste(frame, (0, bar_height))
# Add wrapped text to the bar
text_y_offset = self.CAPTION_PADDING
for line in wrapped_text:
# Calculate text position (centered in the bar horizontally)
if hasattr(draw, "textbbox"):
line_width = (
draw.textbbox((0, 0), line, font=font)[2]
- draw.textbbox((0, 0), line, font=font)[0]
)
line_height = (
draw.textbbox((0, 0), line, font=font)[3]
- draw.textbbox((0, 0), line, font=font)[1]
)
else: # For older Pillow versions, use textsize (deprecated)
line_width, line_height = draw.textsize(line, font=font)
text_x = (new_frame_width - line_width) / 2
draw.text(
(text_x, text_y_offset), line, font=font, fill=self.TEXT_COLOR
)
text_y_offset += line_height
# Reduce colors to optimize GIF and ensure compatibility
new_frame_alpha = new_frame.getchannel("A")
new_frame = new_frame.convert("RGB").convert(
"P", palette=Image.ADAPTIVE, colors=255
)
# If original had transparency, re-apply mask
if gif.info.get("transparency", None) is not None:
new_frame.info["transparency"] = gif.info[
"transparency"
] # Preserve transparency if present
# Masking might be needed here if the original GIF had complex transparency
# For simplicity, we assume simple transparency or opaque.
# If issues arise, more complex alpha compositing might be needed before converting to "P")
frames.append(new_frame)
output_gif_bytes = io.BytesIO()
frames[0].save(
output_gif_bytes,
format="GIF",
save_all=True,
append_images=frames[1:],
duration=gif.info.get(
"duration", self.DEFAULT_GIF_DURATION
), # Use original duration, default to constant
loop=gif.info.get(
"loop", 0
), # Use original loop count, default to infinite
transparency=gif.info.get(
"transparency", None
), # Preserve transparency
disposal=2, # Important for GIFs with transparency and animation
)
output_gif_bytes.seek(0)
return output_gif_bytes
except Exception as e:
print(f"Error in _add_text_to_gif: {e}")
return None
@app_commands.command(
name="captiongif", description="Captions a GIF with the provided text."
)
@app_commands.describe(
caption="The text to add to the GIF.",
url="A URL to a GIF.",
attachment="An uploaded GIF file.",
)
async def caption_gif_slash(
self,
interaction: discord.Interaction,
caption: str,
url: str = None,
attachment: discord.Attachment = None,
):
"""Slash command to caption a GIF."""
await interaction.response.defer(thinking=True)
if not url and not attachment:
await interaction.followup.send(
"You must provide either a GIF URL or attach a GIF file.",
ephemeral=True,
)
return
if url and attachment:
await interaction.followup.send(
"Please provide either a URL or an attachment, not both.",
ephemeral=True,
)
return
image_bytes = None
filename = "captioned_gif.gif"
if url:
if not (
url.startswith("http://tenor.com/")
or url.startswith("https://tenor.com/")
or url.endswith(".gif")
):
await interaction.followup.send(
"The URL must be a direct link to a GIF or a Tenor GIF URL.",
ephemeral=True,
)
return
try:
# Handle Tenor URLs - they often don't directly link to the .gif
# A more robust way is to use Tenor API if available, or try to find the .gif link in the page
# For simplicity, we'll assume if it's a tenor URL, we try to get content and hope it's a GIF
# or that a direct .gif link is provided.
# A common pattern for Tenor is to find a .mp4 or .gif in the HTML if it's a page URL.
# This part might need improvement for robust Tenor URL handling.
# Basic check for direct .gif or try to fetch content
response = requests.get(url, timeout=10)
response.raise_for_status()
content_type = response.headers.get("Content-Type", "").lower()
if "gif" not in content_type and url.endswith(
".gif"
): # If content-type is not gif but url ends with .gif
image_bytes = response.content
elif "gif" in content_type:
image_bytes = response.content
elif "tenor.com" in url: # If it's a tenor URL but not directly a gif
# This is a placeholder for more robust Tenor GIF extraction.
# Often, the actual GIF is embedded. For now, we'll try to fetch and hope.
# A better method would be to parse the HTML for the actual GIF URL.
# Or use the Tenor API if the bot has a key.
# For now, we'll assume the direct URL is good enough or it's a direct .gif from tenor.
# If not, this will likely fail or download HTML.
# A quick hack for some tenor URLs: replace .com/view/ with .com/download/ and hope it gives a direct gif
if "/view/" in url:
potential_gif_url = url.replace(
"/view/", "/download/"
) # This is a guess
# It's better to inspect the page content for the actual media URL
# For now, we'll try the original URL.
pass # Keep original URL for now.
# Attempt to get the GIF from Tenor page (very basic)
if not image_bytes:
page_content = response.text
import re
# Look for a src attribute ending in .gif within an img tag
match = re.search(
r']+src="([^"]+\.gif)"[^>]*>', page_content
)
if match:
gif_url_from_page = match.group(1)
if not gif_url_from_page.startswith(
"http"
): # handle relative URLs if any
from urllib.parse import urljoin
gif_url_from_page = urljoin(url, gif_url_from_page)
response = requests.get(gif_url_from_page, timeout=10)
response.raise_for_status()
if (
"gif"
in response.headers.get("Content-Type", "").lower()
):
image_bytes = response.content
else: # Fallback if no img tag found, try to find a direct media link for tenor
# Tenor often uses a specific div for the main GIF content
# Example: