forked from pckv/pcbot
-
Notifications
You must be signed in to change notification settings - Fork 0
/
image.py
372 lines (290 loc) · 14.1 KB
/
image.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
""" This plugin is designed for various image utility commands.
Commands:
resize """
import logging
import random
import re
from functools import partial
from io import BytesIO
from typing import Any
import discord
from PIL import Image, ImageSequence, ImageOps
import bot
import plugins
from pcbot import utils
url_only = False
# See if we can convert emoji using the emoji.py plugin
try:
from .emoji import get_emote, get_emoji, emote_regex
except:
url_only = True
# See if we can create gifs using imageio
try:
import imageio
except ImportError:
imageio = None
client = plugins.client # type: bot.Client
extension_regex = re.compile(r"image/(?P<ext>\w+)(?:\s|$)")
mention_regex = re.compile(r"<@!?(?P<id>\d+)>")
max_bytes = 4096 ** 2 # 4 MB
max_gif_bytes = 1024 * 6000 # 128kB
def convert_image(image_object: Image.Image, mode: str, real_convert: bool = True):
""" Convert the image object to a specified mode. """
if mode == "RGB" and image_object.mode == "RGBA" and real_convert:
return to_rgb(image_object)
return image_object.convert(mode)
def to_rgb(image_object: Image.Image):
""" Convert to RGB using solution from http://stackoverflow.com/questions/9166400/ """
image_object.load()
background = Image.new("RGB", image_object.size, (0, 0, 0))
background.paste(image_object, mask=image_object.split()[3])
return background
def to_jpg(image_object: Image.Image, quality: int, real_convert=True):
""" Save an image object as JPG and reopen it. """
if image_object.mode != "RGB":
image_object = convert_image(image_object, "RGB", real_convert)
return Image.open(utils.convert_image_object(image_object, "JPEG", quality=quality))
class ImageArg:
def __init__(self, image_object: Image.Image, image_format: str):
self.object = image_object
self.format = image_format
self.extension = image_format.lower()
self.clean_format()
# Figure out if this is a gif by looking for the duration argument. Might only work for gifs
self.gif = bool(self.object.info.get("duration"))
self.gif_bytes = None # For easier upload of gifs, store the bytes in memory
def clean_format(self):
""" Return working options of JPG images. """
if self.extension.lower() == "jpeg":
self.extension = "jpg"
if self.format.lower() == "jpg":
self.format = "JPEG"
def set_extension(self, ext: str):
""" Change the extension of an image. """
self.extension = self.format = ext
def modify(self, function, *args, convert=None, **kwargs):
""" Modify the image object using the given Image function.
This function supplies sequence support. """
if not imageio or not self.gif:
if convert:
self.object = convert_image(self.object, convert)
if isinstance(function, list):
for func in function:
self.object = func(self.object, *args, **kwargs)
else:
self.object = function(self.object, *args, **kwargs)
else:
frames = []
duration = self.object.info.get("duration") / 1000
for frame in ImageSequence.Iterator(self.object):
if convert:
frame = convert_image(frame, convert, real_convert=False)
if isinstance(function, list):
for func in function:
frame = func(frame, *args, **kwargs)
frame_bytes = utils.convert_image_object(frame)
else:
frame_bytes = utils.convert_image_object(function(frame, *args, **kwargs))
frames.append(imageio.v2.imread(frame_bytes, format="PNG"))
# Save the image as bytes and recreate the image object
image_bytes = imageio.mimwrite(imageio.RETURN_BYTES, frames, format=self.format, duration=duration, loop=0)
self.object = Image.open(BytesIO(image_bytes))
self.gif_bytes = image_bytes
async def convert_attachment(attachment: discord.Attachment):
""" Convert an attachment to an image argument.
Returns None if the attachment is not an image.
"""
url = attachment.url
headers = await utils.retrieve_headers(url)
match = extension_regex.search(headers["CONTENT-TYPE"])
if not match:
return None
image_format = match.group("ext")
image_bytes = await utils.download_file(url, bytesio=True)
image_object = Image.open(image_bytes)
return ImageArg(image_object, image_format=image_format)
async def find_prev_image(channel: discord.TextChannel, limit: int = 200):
""" Look for the previous sent image. """
async for message in channel.history(limit=limit):
if message.attachments:
# Try to convert the first attachment
image_arg = await convert_attachment(message.attachments[0])
if not image_arg:
continue
return image_arg
return None
@plugins.argument("{open}url/@user" + ("" if url_only else "/emoji") + "{suffix}{close}", pass_message=True)
async def image(message: discord.Message, url_or_emoji: str):
""" Parse a url, emoji or user mention and return an ImageArg object. """
# Check for local images
if url_or_emoji == ".":
# First see if there is an attachment to this message
image_arg = None
if message.attachments:
image_arg = await convert_attachment(message.attachments[0])
# If there is no attached image, look for an image posted previously
if not image_arg:
image_arg = await find_prev_image(message.channel)
assert image_arg is not None, "Could not find any previously attached image."
return image_arg
# Remove <> if the link looks like a URL, to allow for embed escaped links.
if "http://" in url_or_emoji or "https://" in url_or_emoji:
url_or_emoji = url_or_emoji.strip("<>")
try: # Check if the given string is a url and save the headers for later
headers = await utils.retrieve_headers(url_or_emoji)
except ValueError as e: # Not a valid url, let's see if it's a mention
match = mention_regex.match(url_or_emoji)
if match:
member = message.guild.get_member(int(match.group("id")))
avatar_headers = await utils.retrieve_headers(member.display_avatar.replace(static_format="png").url)
assert not avatar_headers["CONTENT-TYPE"].endswith("gif"), "**GIF avatars are currently unsupported.**"
image_bytes = await utils.download_file(member.display_avatar.replace(size=4096, static_format="png").url,
bytesio=True)
image_object = Image.open(image_bytes)
return ImageArg(image_object, image_format="PNG")
# Nope, not a mention. If we support emoji, we can progress further
assert not url_only, f"`{url_or_emoji}` **is not a valid URL or user mention.**"
# There was no image to get, so let's see if it's an emoji
char = "-".join(hex(ord(c))[2:] for c in url_or_emoji) # Convert to a hex string
image_object = get_emoji(char, size=256)
if image_object:
return ImageArg(image_object, image_format="PNG")
# Not an emoji, perhaps it's an emote
match = emote_regex.match(url_or_emoji)
if match:
image_object = await get_emote(match.group("id"))
if image_object:
return ImageArg(image_object, image_format="PNG")
# Alright, we're out of ideas
raise AssertionError(f"`{url_or_emoji}` **is neither a URL, a mention nor an emoji.**") from e
# The URL was valid so let's make sure it's an image
match = extension_regex.search(headers["CONTENT-TYPE"])
assert match, "**The given URL is not an image.**"
image_format = match.group("ext")
# Make sure the image is not too big
gif = image_format.lower() == "gif"
if "CONTENT-LENGTH" in headers:
size = headers["CONTENT-LENGTH"]
max_size = max_gif_bytes if gif else max_bytes
assert int(size) <= max_size, \
f"**This image exceeds the maximum size of `{max_size // 1024}kB` for this format.**"
elif gif: # If there is no information on the size of the file, we'll refuse if the image is a gif
raise AssertionError("**The size of this GIF is unknown and was therefore rejected.**")
# Download the image and create the object
image_bytes = await utils.download_file(url_or_emoji, bytesio=True)
image_object = Image.open(image_bytes)
return ImageArg(image_object, image_format=image_format)
@plugins.argument("({open}width{close}x{open}height{close} or *{open}scale{close})")
def parse_resolution(res: str):
""" Parse a resolution string.
If the y value is zero, the x value is the number to scale the image with. """
# Check what type of input we're parsing
if res.count("x") == 1:
# Try parsing the numbers in the resolution
x, y = res.split("x")
try:
x = int(x)
y = int(y)
except ValueError as e:
raise AssertionError("**Width or height are not integers.**") from e
# Assign a maximum and minimum size
assert 1 <= x <= 3000 and 1 <= y <= 3000, "**Width and height must be between 1 and 3000.**"
return x, y
if res.startswith("*"):
try:
scale = float(res[1:])
except ValueError as e:
raise AssertionError(rf"**Characters following \* must be a number, not `{res[1:]}`**") from e
# Make sure the scale isn't less than 0. Whatever uses this argument will have to manually check for max size
assert scale > 0, "**Scale must be greater than 0.**"
return scale, 0
return None
def clean_format(image_format: str, extension: str):
""" Return working options of JPG images. """
if extension.lower() == "jpeg":
extension = "jpg"
if image_format.lower() == "jpg":
image_format = "JPEG"
return image_format, extension
async def send_image(message: discord.Message, image_arg: ImageArg, image_format: str = None, **params):
""" Send an image. """
try:
if image_arg.gif and imageio:
image_fp = BytesIO(image_arg.gif_bytes)
else:
image_fp = utils.convert_image_object(image_arg.object, image_format if image_format else image_arg.format,
**params)
except KeyError as e:
await client.send_message(message.channel, f"Image format `{e}` is unsupported.")
else:
await client.send_file(message.channel, image_fp,
filename=f"{message.author.display_name}.{image_arg.extension}")
@plugins.command(pos_check=lambda s: s.startswith("-"))
async def resize(message: discord.Message, image_arg: image, resolution: parse_resolution, *options):
""" Resize an image with the given resolution formatted as `<width>x<height>`
or `*<scale>`. """
# Generate a new image based on the scale
if resolution[1] == 0:
w, h = image_arg.object.size
scale = resolution[0]
assert w * scale < 3000 and h * scale < 3000, "**The result image must be less than 3000 pixels in each axis.**"
resolution = (int(w * scale), int(h * scale))
# Resize and upload the image
image_arg.modify(Image.Image.resize, resolution, Image.NEAREST if "-nearest" in options else Image.LANCZOS,
convert="RGBA")
await send_image(message, image_arg, "PNG")
@plugins.command(pos_check=lambda s: s.startswith("-"), aliases="tilt")
async def rotate(message: discord.Message, image_arg: image, degrees: int, *options, extension: str.lower = None):
""" Rotate an image clockwise using the given degrees. """
if extension:
image_arg.set_extension(extension)
# Rotate and upload the image
image_arg.modify(Image.Image.rotate, -degrees, Image.NEAREST if "-nearest" in options else Image.BICUBIC,
expand=True, convert="RGBA")
await send_image(message, image_arg, "PNG")
@plugins.command(aliases="jpg")
async def jpeg(message: discord.Message, image_arg: image, *effect: utils.choice("meme"),
quality: utils.int_range(f=0, t=100) = 5):
""" Give an image some proper jpeg artifacting.
Valid effects are: `meme` """
image_arg.modify(to_jpg, quality, real_convert="meme" not in effect)
await send_image(message, image_arg)
@plugins.command(aliases="ff")
async def fuckify(message: discord.Message, image_arg: image, seed: Any = None):
""" destroy images """
if seed:
random.seed(seed)
old_size = image_arg.object.size
# Resize to small width and height values
new_size = [random.randint(5, 40) for _ in range(2)]
image_arg.modify([
partial(Image.Image.resize, size=new_size, resample=Image.LANCZOS),
partial(to_jpg, quality=random.randint(3, 30)),
partial(Image.Image.resize, size=old_size, resample=Image.LANCZOS),
partial(to_jpg, quality=random.randint(1, 20)),
], convert="RGBA")
await send_image(message, image_arg, "PNG")
@plugins.command()
async def invert(message: discord.Message, image_arg: image):
""" Invert the colors of an image. """
image_arg.modify(ImageOps.invert, convert="RGB")
await send_image(message, image_arg, quality=100)
@plugins.command()
async def flip(message: discord.Message, image_arg: image, extension: str.lower = None):
""" Flip an image in the y-axis. """
if extension:
image_arg.set_extension(extension)
# Flip the image
image_arg.modify(Image.Image.transpose, Image.FLIP_TOP_BOTTOM)
try:
await send_image(message, image_arg)
except IOError:
await client.say(message, "**The image format is not supported (must be L or RGB)**")
@plugins.command()
async def mirror(message: discord.Message, image_arg: image, extension: str.lower = None):
""" Mirror an image along the x-axis. """
if extension:
image_arg.set_extension(extension)
# Mirror the image
image_arg.modify(Image.Image.transpose, Image.FLIP_LEFT_RIGHT)
await send_image(message, image_arg)