253 lines
9.6 KiB
Python
253 lines
9.6 KiB
Python
import discord
|
|
from discord.ext import commands
|
|
from discord import app_commands
|
|
from PIL import Image, ImageSequence, UnidentifiedImageError
|
|
import os
|
|
import io
|
|
import tempfile
|
|
import traceback
|
|
|
|
|
|
class GIFOptimizerCog(commands.Cog):
|
|
def __init__(self, bot):
|
|
self.bot = bot
|
|
|
|
async def _optimize_gif_internal(
|
|
self,
|
|
input_bytes: bytes,
|
|
colors: int,
|
|
dither_on: bool,
|
|
crop_box_str: str,
|
|
resize_dimensions_str: str,
|
|
):
|
|
"""
|
|
Internal function to optimize a GIF from bytes, returning optimized bytes.
|
|
Handles file I/O using temporary files.
|
|
"""
|
|
input_path = None
|
|
output_path = None
|
|
try:
|
|
# Create a temporary input file
|
|
with tempfile.NamedTemporaryFile(
|
|
delete=False, suffix=".gif"
|
|
) as temp_input_file:
|
|
temp_input_file.write(input_bytes)
|
|
input_path = temp_input_file.name
|
|
|
|
# Create a temporary output file
|
|
with tempfile.NamedTemporaryFile(
|
|
delete=False, suffix=".gif"
|
|
) as temp_output_file:
|
|
output_path = temp_output_file.name
|
|
|
|
# Parse crop and resize arguments
|
|
crop_box_tuple = None
|
|
if crop_box_str:
|
|
try:
|
|
parts = [int(p.strip()) for p in crop_box_str.split(",")]
|
|
if len(parts) == 4:
|
|
crop_box_tuple = tuple(parts)
|
|
else:
|
|
raise ValueError(
|
|
"Crop argument must be four integers: left,top,right,bottom (e.g., '10,20,100,150')"
|
|
)
|
|
except ValueError as e:
|
|
return None, f"Invalid crop format: {e}"
|
|
|
|
resize_dims_tuple = None
|
|
if resize_dimensions_str:
|
|
try:
|
|
parts_str = resize_dimensions_str.replace("x", ",").split(",")
|
|
parts = [int(p.strip()) for p in parts_str]
|
|
if len(parts) == 2:
|
|
if parts[0] > 0 and parts[1] > 0:
|
|
resize_dims_tuple = tuple(parts)
|
|
else:
|
|
raise ValueError(
|
|
"Resize dimensions (width, height) must be positive integers."
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
"Resize argument must be two positive integers: width,height (e.g., '128,128' or '128x128')"
|
|
)
|
|
except ValueError as e:
|
|
return None, f"Invalid resize format: {e}"
|
|
|
|
dither_method = Image.Dither.FLOYDSTEINBERG if dither_on else None
|
|
num_colors = max(2, min(colors, 256))
|
|
|
|
# --- Original optimize_gif logic, adapted ---
|
|
img = Image.open(input_path)
|
|
|
|
original_loop = img.info.get("loop", 0)
|
|
original_transparency = img.info.get("transparency")
|
|
|
|
processed_frames = []
|
|
durations = []
|
|
disposals = []
|
|
|
|
for i, frame_image in enumerate(ImageSequence.Iterator(img)):
|
|
durations.append(frame_image.info.get("duration", 100))
|
|
disposals.append(frame_image.info.get("disposal", 2))
|
|
|
|
current_frame = frame_image.copy()
|
|
|
|
if crop_box_tuple:
|
|
try:
|
|
current_frame = current_frame.crop(crop_box_tuple)
|
|
except Exception as crop_error:
|
|
# Log warning, but don't fail the entire process for a single frame crop error
|
|
print(
|
|
f"Warning: Could not apply crop box {crop_box_tuple} to frame {i+1}. Error: {crop_error}"
|
|
)
|
|
# Optionally, you could decide to skip this frame or use the original frame if cropping is critical.
|
|
# For now, we'll let the error propagate if it's a critical image error, otherwise proceed.
|
|
|
|
if resize_dims_tuple:
|
|
try:
|
|
# Use Image.LANCZOS directly
|
|
current_frame = current_frame.resize(
|
|
resize_dims_tuple, Image.LANCZOS
|
|
)
|
|
except Exception as resize_error:
|
|
print(
|
|
f"Warning: Could not resize frame {i+1} to {resize_dims_tuple}. Error: {resize_error}"
|
|
)
|
|
|
|
frame_rgba = current_frame.convert("RGBA")
|
|
quantized_frame = frame_rgba.convert(
|
|
"P",
|
|
palette=Image.Palette.ADAPTIVE,
|
|
colors=num_colors,
|
|
dither=dither_method,
|
|
)
|
|
processed_frames.append(quantized_frame)
|
|
|
|
if not processed_frames:
|
|
return None, "No frames processed from the GIF."
|
|
|
|
processed_frames[0].save(
|
|
output_path,
|
|
save_all=True,
|
|
append_images=processed_frames[1:],
|
|
optimize=True,
|
|
duration=durations,
|
|
loop=original_loop,
|
|
disposal=disposals,
|
|
transparency=original_transparency,
|
|
)
|
|
|
|
with open(output_path, "rb") as f:
|
|
optimized_bytes = f.read()
|
|
|
|
input_size = len(input_bytes)
|
|
output_size = len(optimized_bytes)
|
|
reduction_percentage = (
|
|
(input_size - output_size) / input_size * 100 if input_size > 0 else 0
|
|
)
|
|
|
|
stats = (
|
|
f"Original size: {input_size / 1024:.2f} KB\n"
|
|
f"Optimized size: {output_size / 1024:.2f} KB\n"
|
|
f"Reduction: {reduction_percentage:.2f}%"
|
|
)
|
|
return optimized_bytes, stats
|
|
|
|
except FileNotFoundError:
|
|
return None, "Internal error: Temporary file not found."
|
|
except UnidentifiedImageError:
|
|
return (
|
|
None,
|
|
"Cannot identify image file. It might be corrupted or not a supported GIF format.",
|
|
)
|
|
except Exception as e:
|
|
traceback.print_exc() # Print full traceback to console for debugging
|
|
return None, f"An unexpected error occurred during GIF optimization: {e}"
|
|
finally:
|
|
# Clean up temporary files
|
|
if input_path and os.path.exists(input_path):
|
|
os.remove(input_path)
|
|
if output_path and os.path.exists(output_path):
|
|
os.remove(output_path)
|
|
|
|
@commands.command(name="optimizegif", description="Optimizes a GIF attachment.")
|
|
async def optimize_gif_prefix(
|
|
self,
|
|
ctx: commands.Context,
|
|
attachment: discord.Attachment,
|
|
colors: int = 128,
|
|
dither: bool = True,
|
|
crop: str = None,
|
|
resize: str = None,
|
|
):
|
|
if not attachment.filename.lower().endswith(".gif"):
|
|
await ctx.send("Please provide a GIF file.")
|
|
return
|
|
|
|
await ctx.defer()
|
|
try:
|
|
input_bytes = await attachment.read()
|
|
optimized_bytes, stats = await self._optimize_gif_internal(
|
|
input_bytes, colors, dither, crop, resize
|
|
)
|
|
|
|
if optimized_bytes:
|
|
file = discord.File(
|
|
io.BytesIO(optimized_bytes),
|
|
filename=f"optimized_{attachment.filename}",
|
|
)
|
|
await ctx.send(f"GIF optimized successfully!\n{stats}", file=file)
|
|
else:
|
|
await ctx.send(f"Failed to optimize GIF: {stats}")
|
|
except Exception as e:
|
|
await ctx.send(f"An error occurred: {e}")
|
|
traceback.print_exc()
|
|
|
|
@app_commands.command(name="optimizegif", description="Optimizes a GIF attachment.")
|
|
@app_commands.describe(
|
|
attachment="The GIF file to optimize.",
|
|
colors="Number of colors to reduce to (e.g., 256, 128, 64). Max 256.",
|
|
dither="Enable Floyd-Steinberg dithering (improves quality, slightly slower).",
|
|
crop="Crop the GIF. Provide as 'left,top,right,bottom' (e.g., '10,20,100,150').",
|
|
resize="Resize the GIF. Provide as 'width,height' (e.g., '128,128' or '128x128').",
|
|
)
|
|
async def optimize_gif_slash(
|
|
self,
|
|
interaction: discord.Interaction,
|
|
attachment: discord.Attachment,
|
|
colors: app_commands.Range[int, 2, 256] = 128,
|
|
dither: bool = True,
|
|
crop: str = None,
|
|
resize: str = None,
|
|
):
|
|
if not attachment.filename.lower().endswith(".gif"):
|
|
await interaction.response.send_message(
|
|
"Please provide a GIF file.", ephemeral=True
|
|
)
|
|
return
|
|
|
|
await interaction.response.defer()
|
|
try:
|
|
input_bytes = await attachment.read()
|
|
optimized_bytes, stats = await self._optimize_gif_internal(
|
|
input_bytes, colors, dither, crop, resize
|
|
)
|
|
|
|
if optimized_bytes:
|
|
file = discord.File(
|
|
io.BytesIO(optimized_bytes),
|
|
filename=f"optimized_{attachment.filename}",
|
|
)
|
|
await interaction.followup.send(
|
|
f"GIF optimized successfully!\n{stats}", file=file
|
|
)
|
|
else:
|
|
await interaction.followup.send(f"Failed to optimize GIF: {stats}")
|
|
except Exception as e:
|
|
await interaction.followup.send(f"An error occurred: {e}")
|
|
traceback.print_exc()
|
|
|
|
|
|
async def setup(bot):
|
|
await bot.add_cog(GIFOptimizerCog(bot))
|