nikita
2867582c01
- Updated README.md - Updated examples - `MatrixBot` does not redefine methods from classes it is composed of anymore
129 lines
3.6 KiB
Python
129 lines
3.6 KiB
Python
"""
|
|
This example implements Matrix bot that generates a pixelized noise image with
|
|
specified maximum R, G and B values.
|
|
|
|
It uses environment variables to specify authorization data. Use Ctrl+C to stop
|
|
the bot.
|
|
"""
|
|
|
|
import asyncio
|
|
from io import BytesIO
|
|
import os
|
|
import time
|
|
import logging
|
|
import random
|
|
from PIL import Image, ImageFilter
|
|
|
|
from mab import *
|
|
|
|
from _environment import check_environment
|
|
|
|
async def on_gen_command(data: EventContext) -> None:
|
|
"""This callback is called when `!gen R G B` command is received."""
|
|
# convert R, G and B to floats
|
|
try:
|
|
r, g, b = [float(v) for v in data[CTX_CMD_ARGS]]
|
|
except:
|
|
await data.bot.send.text(data.room, "Invalid arguments")
|
|
return
|
|
await data.bot.send.text(data.room, "Generating the noise...")
|
|
# create the basic noise
|
|
img = Image.new("RGB", (16, 16))
|
|
for x in range(img.width):
|
|
for y in range(img.height):
|
|
col = (random.random() * r, random.random() * g, random.random() * b)
|
|
img.putpixel(
|
|
(x, y),
|
|
tuple(int(c * 255) for c in col)
|
|
)
|
|
# pixelized upscale
|
|
img = img.resize((2048, 2048), resample=Image.Resampling.NEAREST)
|
|
# save to buffer
|
|
buf = BytesIO()
|
|
img.save(buf, format="PNG")
|
|
buf.seek(0)
|
|
buf = buf.read()
|
|
# send
|
|
await data.bot.send.image_bytes(data.room, buf, "noise.png")
|
|
|
|
async def on_image(ctx: EventContext) -> None:
|
|
"""This callback is called when an image is received."""
|
|
# download
|
|
s = time.time()
|
|
data = await ctx.bot.download.file(ctx)
|
|
took_time = time.time() - s
|
|
# convert to Image
|
|
with BytesIO(data) as buf:
|
|
img = Image.open(buf)
|
|
# apply effects
|
|
blur = ImageFilter.GaussianBlur(
|
|
radius=min(img.size[0] // 10, 5)
|
|
)
|
|
img = img.filter(blur)
|
|
# save to buffer
|
|
buf = BytesIO()
|
|
img.save(buf, format="PNG")
|
|
buf.seek(0)
|
|
buf = buf.read()
|
|
# send
|
|
await ctx.bot.send.image_bytes(
|
|
ctx.room,
|
|
buf,
|
|
"blurred.png",
|
|
text="Download and decryption took %.4f seconds" % took_time
|
|
)
|
|
|
|
async def on_wrong_message(data: EventContext) -> None:
|
|
"""This callback is called when a wrong message is received."""
|
|
await data.bot.send.text(
|
|
data.room,
|
|
"Text me something like <code>!gen 0.1 0.7 1.0</code> or send an image to blur"
|
|
)
|
|
|
|
async def main() -> None:
|
|
"""Application entry point"""
|
|
logging.basicConfig(level=logging.INFO)
|
|
logging.getLogger("nio").setLevel(logging.CRITICAL+1)
|
|
check_environment()
|
|
|
|
config = MatrixBotConfig(
|
|
matrix_homeserver_url=os.environ["MATRIX_HOMESERVER"],
|
|
matrix_username_localpart=os.environ["MATRIX_USERNAME"],
|
|
storage_directory="session_storage"
|
|
)
|
|
bot = MatrixBot(config)
|
|
command_filter = MessageTypeFilter(MessageType.TEXT) & BodyCommandFilter(
|
|
verbs=["gen"],
|
|
min_args=3,
|
|
max_args=3
|
|
)
|
|
|
|
# callback for message that
|
|
# 1. are sent not by this bot
|
|
# 2. do match the command filter
|
|
bot.add_callback(
|
|
~SenderIsBotFilter() & command_filter,
|
|
on_gen_command)
|
|
|
|
# callback for new image message
|
|
bot.add_callback(
|
|
~SenderIsBotFilter() & MessageTypeFilter(MessageType.IMAGE) & NewMessageFilter(),
|
|
on_image
|
|
)
|
|
|
|
# callback for message that
|
|
# 1. are sent not by this bot
|
|
# 2. are new messages (not edits)
|
|
bot.add_callback(
|
|
~SenderIsBotFilter() & BodyExistsFilter() & NewMessageFilter(),
|
|
on_wrong_message
|
|
)
|
|
|
|
# run until Ctrl+C
|
|
try:
|
|
await bot.run()
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main()) |