156 lines
4.9 KiB
Python
156 lines
4.9 KiB
Python
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())
|