Disocrd.pyを用いたメインカラー抽出botの作成

Disocrd.pyを使用してDiscordにアップロードされた画像からメインカラーを抽出するbotを作成する。

方法

k-means法でクラスタリングを行います

ピクセルごとに指定されたクラスタ数でクラスタリングを行い、書くクラスタの中央の座標の色をメインカラーとします

Discord.pyのほうでは、画像と一緒に送られてくる文字が/picのときに反応するようにしています。

使用するmodule

pip

Pillow
OpenCV
sklearn.cluster
numpy
random
request
io

コード解説

import文

from PIL import Image, ImageDraw, ImageFont
import cv2
from sklearn.cluster import KMeans
import numpy as np
from numpy import linalg as LA
import random
import discord
import requests
import io

import文の解説はは省略します。

discord.pyの設定

color_stripe.py

client = discord.Client()
Token = 'Enter Your TOKEN'

Enter Your TOKENにはあなたのbotのtokenを入力してください。

botの反応の設定

color_stripe.py

@client.event
async def on_message(message):
    if message.content.startswith('/pic'):
        print(message.attachments[0].url)
        r = requests.get(message.attachments[0].url)
        img = Image.open(io.BytesIO(r.content))
        img.save("image.png")
        color_arr = extract_main_color(img_path, 7)
        show_tiled_main_color(color_arr)
        draw_random_stripe(color_arr, img_path)
        file = discord.File("./image/stripe_image.png", filename="stripe.png")
        await message.channel.send(file=file)

一つずつ説明すると

color_stripe.py

async def on_message(message):
    if message.content.startswith('/pic'):

メッセージを受け取り、それが/picだったら処理を続行する

color_stripe.py

print(message.attachments[0].url)

送られてきた画像のurlを出力します

color_stripe.py

r = requests.get(message.attachments[0].url)
        img = Image.open(io.BytesIO(r.content))
        img.save("image.png")

requestを使用して画像を保存します。

color_stripe.py

color_arr = extract_main_color(img_path, 7)
        show_tiled_main_color(color_arr)
        draw_random_stripe(color_arr, img_path)
        file = discord.File("./image/stripe_image.png", filename="stripe.png")
        await message.channel.send(file=file)

下に記述するdef文の処理を呼び出し、Discordに送信します。
今回はshow_tiled_main_colorのみを送信します。
def文中について説明します

color_extra.py

def draw_random_stripe(color_arr, img_path):
    width = 1024
    height = 1024

    stripe_color_img = Image.new(
        mode='RGB', size=(width, height), color='#333333')
    current_height = 0
    while current_height < height:
        random_index = random.randrange(color_arr.shape[0])
        color_hex_str = '#%02x%02x%02x' % tuple(color_arr[random_index])
        random_height = random.randrange(5, 70)
        color_img = Image.new(
            mode='RGB', size=(width, random_height),
            color=color_hex_str)
        stripe_color_img.paste(
            im=color_img,
            box=(0, current_height))
        current_height += random_height
    stripe_color_img.show()
    #stripe_color_img.save('./image/stripe_' + img_path)

メインカラーをランダムに選択しランダムな縦幅で矩形描画、これを繰り返しストライプを作成し，表示し保存します。

color_extra.py

def show_tiled_main_color(color_arr):

    IMG_SIZE = 64
    MARGIN = 15
    width = IMG_SIZE * color_arr.shape[0] + MARGIN * 2
    height = IMG_SIZE + MARGIN * 2

    tiled_color_img = Image.new(
        mode='RGB', size=(width, height), color='#333333')

    for i, rgb_arr in enumerate(color_arr):
        color_hex_str = '#%02x%02x%02x' % tuple(rgb_arr)
        color_img = Image.new(
            mode='RGB', size=(IMG_SIZE, IMG_SIZE),
            color=color_hex_str)
        tiled_color_img.paste(
            im=color_img,
            box=(MARGIN + IMG_SIZE * i, MARGIN))

    tiled_color_img.show()
    tiled_color_img.save('./image/stripe_' + img_path)

抽出したメインカラーを横並びにPILを使用して作成します。

多分最後が一番大事です(笑)

color_extra.py

def extract_main_color(img_path, color_num):

    cv2_img = cv2.imread(img_path)
    cv2_img = cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB)
    cv2_img = cv2_img.reshape(
        (cv2_img.shape[0] * cv2_img.shape[1], 3))

    cluster = KMeans(n_clusters=color_num)
    cluster.fit(X=cv2_img)
    cluster_centers_arr = cluster.cluster_centers_.astype(
        int, copy=False)
    trans_color = cv2_img[0]
    cluster_centers_arr = np.array([i for i in cluster_centers_arr if LA.norm(np.array(i - trans_color), 2) > 50])
    print("extracted colors array:")
    print(cluster_centers_arr)
    return cluster_centers_arr

画像のメインカラーを抽出します。
KMeans方を使用して行っています。
また

color_extra.py

trans_color = cv2_img[0]

にて左上の色は透過色とみなす(背景が白等の場合、それを無視するため)
また、透過色に近い色を無視することもできます。

color_extra.py

cluster_centers_arr = np.array([i for i in cluster_centers_arr if LA.norm(np.array(i - trans_color), 2) > 50])

これでプログラムのほとんどはできました。
最後にdiscordから送られてきた画像を参照、あとDiscordのBotのToeknを設定します。

img_path = 'image.png'
client.run('TOKEN')

これにて完成です。お疲れさまでした！

一応github上にファイルを上げたので参考にしてください！
https://github.com/Tps-F/color_extract_discordbot/blob/main/color_picker.py

参考情報