From db6918fa6cce88fd9d4d55654f1436aafdee9685 Mon Sep 17 00:00:00 2001 From: Nathan Rashleigh Date: Sat, 6 Apr 2024 16:54:58 +1100 Subject: [PATCH] meili --- .gitignore | 7 +- TODO.md | 7 ++ image.py | 26 +++++++ main.py | 155 +++++++++++++++++++++++++++++++++++++ memesearch.sublime-project | 18 +++++ requirements.txt | 3 +- 6 files changed, 213 insertions(+), 3 deletions(-) create mode 100644 TODO.md create mode 100644 image.py create mode 100644 main.py create mode 100644 memesearch.sublime-project diff --git a/.gitignore b/.gitignore index 173acff..489dad4 100644 --- a/.gitignore +++ b/.gitignore @@ -1,9 +1,12 @@ .env -venv/ +.venv/ +gimmeme-d052de2b9472.json __pycache__/ # drop the workspace file cause it leaks info about personal machines *.sublime-workspace -data/ \ No newline at end of file +data/ +meili_data/ +data.bak/ \ No newline at end of file diff --git a/TODO.md b/TODO.md new file mode 100644 index 0000000..5d7f3de --- /dev/null +++ b/TODO.md @@ -0,0 +1,7 @@ +# MemeSearch + +## TODO + +[ ] Backfill last N images from a channel +[ ] Set oldest_img_id for channel +[ ] diff --git a/image.py b/image.py new file mode 100644 index 0000000..7d446ed --- /dev/null +++ b/image.py @@ -0,0 +1,26 @@ +from google.cloud import vision +from google.cloud.vision_v1 import types + +from dotenv import load_dotenv + +load_dotenv() + +crows = 'https://cdn.discordapp.com/attachments/367936426962190336/1222335573855834214/AjPuqci.png?ex=6615d7b6&is=660362b6&hm=e364398ce7401334298eb82bddedab56d944f5833508fde4b07363b329a95c6f&' + +# Authenticate the client +client = vision.ImageAnnotatorClient() + +# Perform text detection +features = [types.Feature(type=types.Feature.Type.TEXT_DETECTION)] +request = { + 'image': { + 'source': { 'image_uri': crows } + }, + 'features': features +} + +response: types.AnnotateImageResponse = client.annotate_image(request=request) +texts = response.text_annotations[0].description + +print(texts) + diff --git a/main.py b/main.py new file mode 100644 index 0000000..9ba5631 --- /dev/null +++ b/main.py @@ -0,0 +1,155 @@ +import os +import asyncio +import aiohttp +import json +import meilisearch + +from dataclasses import dataclass +from typing import Dict, Any + +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 + +# need to change the env for the google api calls +load_dotenv() +config = dotenv_values(".env") + +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 MemeMessage: + id: int + url: str + guild_name: str + channel_name: str + + text: str | None + vision_json: 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] + + async def save_img(self, session: aiohttp.ClientSession): + output_dir = self.output_dir() + + # 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(output_dir, exist_ok=True); + + async with session.get(self.url) as response: + if response.status == 200: + with open(f'{output_dir}/{self.id}{ext}', 'wb') as f: + f.write(await response.read()) + + def output_dir(self): + return f'./data/{self.guild_name}/{self.channel_name}' + + 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_meili_document(self): + return { + 'id': self.id, + 'text': self.text, + 'url': self.url, + } + +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('http://localhost:7700', 'aSampleMasterKey') + documents = [meme.to_meili_document() for meme in memes] + client.index('memes').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()) diff --git a/memesearch.sublime-project b/memesearch.sublime-project new file mode 100644 index 0000000..a4f79a1 --- /dev/null +++ b/memesearch.sublime-project @@ -0,0 +1,18 @@ +{ + "folders": + [ + { + "path": "." + } + ], + "build_systems": [ + { + "name": "Python - memesearch", + "cmd": ["$project_path/.venv/bin/python3", "-u", "main.py"] + }, + { + "name": "Python - memesearch - current file", + "cmd": ["$project_path/.venv/bin/python3", "-u", "$file"] + } + ] +} diff --git a/requirements.txt b/requirements.txt index f06d7c4..b19a9aa 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ python-dotenv aiohttp -google-cloud-vision \ No newline at end of file +google-cloud-vision +meilisearch \ No newline at end of file