memestream/main.py

import os
import asyncio
import aiohttp
import json
import meilisearch

from io import BytesIO
from dataclasses import dataclass
from typing import Dict, Any, List

from dotenv import dotenv_values, load_dotenv
from google.cloud import vision as google_vision
from google.protobuf.json_format import MessageToDict, MessageToJson
from google.cloud.vision_v1 import types
from PIL import Image

# need to change the env for the google api calls
load_dotenv()
config = dotenv_values(".env")

DATA_DIR = './data/memesearch'
MEILI_INDEX = 'memes'
MEILI_URL = 'http://localhost:7700'

CHANNEL_IDS = [
    367936426962190336, # TPM shitposts
    1095060717117374687, # aella memes
]

headers = {
    'Authorization': config['DISCORD_AUTH'],
}

# TODO
# [ ] Backfill last N images from a channel
# [ ] Set oldest_img_id for channel
# [ ] save img as json + imagefile
# [ ] check extensions (ignore mp4 maybe)

@dataclass
class MemeDocument:
    id: int
    url: str
    text: str | None
    tags: List[str]

@dataclass
class MemeMessage:
    id: int
    url: str
    guild_name: str
    channel_name: str

    vision_json: str | None
    text: str | None

    # maybe don't need these
    # kind of extra
    message_json: Dict[str, Any]
    guild_json: Dict[str, Any]
    channel_json: Dict[str, Any]

    def output_dir(self):
        return f'{DATA_DIR}/guilds/{self.guild_name}/{self.channel_name}'

    def img_dir(self):
        return f"{DATA_DIR}/imgs/"

    async def save_img(self, session: aiohttp.ClientSession):
        webp_output = f'{DATA_DIR}/imgs/'

        # there's a big param of some kind after the image url that we remove
        url_filename = os.path.basename(self.url).split('?')[0]
        (filename, ext) = os.path.splitext(url_filename)
        print(f"filename={filename} ext={ext}")

        os.makedirs(self.output_dir(), exist_ok=True);
        os.makedirs(self.img_dir(), exist_ok=True);

        async with session.get(self.url) as response:
            if response.status == 200:
                res_bytes = await response.read()

                with open(f'{self.output_dir()}/{self.id}{ext}', 'wb') as f:
                    f.write(res_bytes)

                convert_extensions = ['jpg', 'png', 'jpeg', 'webp']
                if ext.lower() in convert_extensions:
                    image = Image.open(BytesIO(res_bytes))
                    image.save(f"{self.img_dir()}/{self.id}.webp", format='WebP')

    def annotate_image_request(self, vision):
        request = {
            'image': { 'source': { 'image_uri': self.url } },
            'features': [types.Feature(type=types.Feature.Type.TEXT_DETECTION)],
        }

        response: types.AnnotateImageResponse = vision.annotate_image(request=request)
        self.vision_json = json.loads(types.AnnotateImageResponse.to_json(response))

        if len(response.text_annotations) > 0:
            self.text = response.text_annotations[0].description

        return request

    def persist(self):
        with open(f'{self.output_dir()}/{self.id}.json', 'w') as f:
            json.dump(vars(self), f, indent=4)

    def to_document(self):
        return MemeDocument(
            id=self.id,
            url=f'{DATA_DIR}/{self.id}.webp',
            text=self.text,
            tags=[self.guild_name, self.channel_name],
        )

def filter_message_img_urls(messages):
    url_lists = map(
        lambda msg: list(map(lambda attachment: (attachment['url'], msg), msg['attachments']))
    , messages)
    return [item for group in url_lists for item in group]

async def download_channel_imgs(session, vision, channel_id):
    async def fetch_channel_messages(channel_id):
        url = f"https://discord.com/api/v9/channels/{channel_id}/messages"
        async with session.get(url, headers=headers) as res:
            return await res.json()

    async def fetch_guild_details(guild_id):
        url = f"https://discord.com/api/v9/guilds/{guild_id}"
        async with session.get(url, headers=headers) as res:
            data = await res.json()

            keys_to_remove = ['emojis', 'stickers', 'roles']
            [data.pop(key, None) for key in keys_to_remove]

            return data

    async def fetch_channel(channel_id):
        url = f"https://discord.com/api/v9/channels/{channel_id}"
        async with session.get(url, headers=headers) as res:
            return await res.json()

    channel = await fetch_channel(channel_id)
    guild = await fetch_guild_details(channel['guild_id'])
    print(f"Fetching messages for channel guild={guild['name']} name={channel['name']} id={channel['id']}")
    messages = await fetch_channel_messages(channel_id)
    urls = filter_message_img_urls(messages)

    memes = list(map(
        lambda msg: MemeMessage(
            id=msg[1]['id'],
            url=msg[0],
            guild_name=guild['name'],
            channel_name=channel['name'],
            guild_json=guild,
            channel_json=channel,
            message_json=msg[1],
            vision_json=None,
            text=None,
        )
    , urls))

    tasks = [meme.save_img(session) for meme in memes]
    await asyncio.gather(*tasks)

    [meme.annotate_image_request(vision) for meme in memes]
    [meme.persist() for meme in memes]

    client = meilisearch.Client(MEILI_URL, config['MEILI_MASTER_KEY'])

    # have to vars the dataclass out for meili
    documents = [vars(meme.to_document()) for meme in memes]
    client.index(MEILI_INDEX).add_documents(documents)

async def main():
    vision = google_vision.ImageAnnotatorClient()
    async with aiohttp.ClientSession() as session:
        channels = [download_channel_imgs(session, vision, channel_id) for channel_id in CHANNEL_IDS]
        memes = await asyncio.gather(*channels)

asyncio.run(main())