236 lines
7.8 KiB
Python
236 lines
7.8 KiB
Python
import os
|
|
import asyncio
|
|
import json
|
|
import itertools
|
|
|
|
from urllib3.util import parse_url
|
|
from pathlib import Path
|
|
from io import BytesIO
|
|
from dataclasses import dataclass
|
|
from uuid import uuid4, UUID
|
|
|
|
import aiohttp
|
|
|
|
from aiobotocore.session import get_session as aiobotocore_get_session
|
|
from types_aiobotocore_s3.client import S3Client
|
|
|
|
from PIL import Image
|
|
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
|
|
import meilisearch
|
|
from urllib3.util import Url, url
|
|
|
|
from sqlalchemy import create_engine
|
|
from db.migrate import migrate_db
|
|
from db.attachment import Attachment
|
|
from db.meme import Meme
|
|
|
|
# need to change the env for the google api calls
|
|
load_dotenv()
|
|
config = dotenv_values(".env")
|
|
|
|
DATA_DIR = './data/memestream'
|
|
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: UUID
|
|
url: Url
|
|
text: str
|
|
tags: list[str]
|
|
|
|
def detect_text(url, vision):
|
|
request = {
|
|
'image': { 'source': { 'image_uri': url } },
|
|
'features': [types.Feature(type=types.Feature.Type.TEXT_DETECTION)],
|
|
}
|
|
|
|
response: types.AnnotateImageResponse = vision.annotate_image(request=request)
|
|
vision_json = json.loads(types.AnnotateImageResponse.to_json(response))
|
|
|
|
return request
|
|
|
|
def filter_message_img_urls(messages):
|
|
url_lists = map(
|
|
lambda msg: list(map(lambda attachment: (parse_url(attachment['url']), msg), msg['attachments']))
|
|
, messages)
|
|
return [item for group in url_lists for item in group]
|
|
|
|
async def download_channel_attachments(session, storage, 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()
|
|
|
|
|
|
async def upload_object(object_name, body):
|
|
storage_bucket = config['STORAGE_BUCKET']
|
|
if not storage_bucket:
|
|
raise RuntimeError(f"config missing STORAGE_BUCKET")
|
|
return await storage.put_object(Bucket=storage_bucket, Key=object_name, Body=body)
|
|
|
|
async def download_attachment(
|
|
session: aiohttp.ClientSession,
|
|
storage: S3Client,
|
|
attachment: Attachment
|
|
):
|
|
async with session.get(attachment.url) as response:
|
|
if response.status == 200:
|
|
# make a webp conversion for obvious image types
|
|
res_bytes = await response.read()
|
|
|
|
(filename, ext) = attachment.file_details()
|
|
object_name = attachment.object_name()
|
|
print(f"uploading object_name={object_name} filename={filename} ext={ext}")
|
|
|
|
if attachment.convert_to_webp():
|
|
image = Image.open(BytesIO(res_bytes))
|
|
webp_output = BytesIO()
|
|
image.save(webp_output, format='WebP')
|
|
webp_output.seek(0)
|
|
return await upload_object(object_name=object_name, body=webp_output)
|
|
|
|
return await upload_object(object_name=object_name, body=res_bytes)
|
|
|
|
def create_attachments(message, channel, guild):
|
|
urls = [parse_url(attachment['url']) for attachment in message['attachments']]
|
|
|
|
def generate_attachment(url: Url):
|
|
if not url.path:
|
|
return None
|
|
|
|
(filename, ext) = os.path.splitext(os.path.basename(url.path))
|
|
extensions = ['jpg', 'png', 'jpeg', 'webp', 'mov', 'gif', 'mp4']
|
|
if ext.lower()[1:] not in extensions:
|
|
return None
|
|
|
|
return Attachment(
|
|
id=uuid4(),
|
|
msg_id=int(message['id']),
|
|
message=message,
|
|
guild=guild,
|
|
url=url.url,
|
|
channel=channel
|
|
)
|
|
|
|
return [generate_attachment(url) for url in urls]
|
|
|
|
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)
|
|
|
|
# at this point we have the channel messages, the guild (server) and channel info
|
|
|
|
# first thing we have to do is parse out the urls in the msg attachments
|
|
# once we have those we can generate where we should put them in our storage
|
|
# should only persist them later once the processing is done
|
|
attachment_groups = [create_attachments(message, channel, guild) for message in messages]
|
|
|
|
# super wack flatmap and filter
|
|
attachments = list(filter(None, itertools.chain.from_iterable(attachment_groups)))
|
|
|
|
# NOW
|
|
# -------------------
|
|
# insert attachment
|
|
# upload to s3
|
|
# fetch vision for relevant attachments
|
|
# create meme from attachment and vision
|
|
# insert meme
|
|
# add documents to meili
|
|
|
|
# tasks = [download_attachment(session, storage, url, message) for (url, message) in urls_and_messages]
|
|
# await asyncio.gather(*tasks)
|
|
|
|
# [attachment.detect_text(vision) for attachment in attachments]
|
|
# [attachment.persist() for attachment in attachments]
|
|
|
|
# client = meilisearch.Client(MEILI_URL, config['MEILI_MASTER_KEY'])
|
|
|
|
# # have to vars the dataclass out for meili
|
|
# documents = [vars(attachment.to_meme()) for attachment in attachments]
|
|
# client.index(MEILI_INDEX).add_documents(documents)
|
|
|
|
async def check_bucket_exists(storage, bucket_name):
|
|
try:
|
|
await storage.head_bucket(Bucket=bucket_name)
|
|
return True
|
|
except storage.exceptions.ClientError as e:
|
|
error_code = e.response.get('Error', {}).get('Code')
|
|
if error_code == '404':
|
|
return False # Bucket does not exist
|
|
else:
|
|
raise # Reraise other exceptions
|
|
|
|
async def main():
|
|
assert config['STORAGE_URL']
|
|
assert config['POSTGRES_URL']
|
|
|
|
vision = google_vision.ImageAnnotatorClient()
|
|
|
|
# db setup
|
|
db_url = config['POSTGRES_URL']
|
|
db_engine = create_engine(db_url, echo=True)
|
|
migrate_db(db_engine, db_url)
|
|
|
|
aioboto = aiobotocore_get_session()
|
|
async with aioboto.create_client(
|
|
's3',
|
|
endpoint_url=config['STORAGE_URL'],
|
|
aws_access_key_id=config['STORAGE_ACCESS_KEY'],
|
|
aws_secret_access_key=config['STORAGE_SECRET_KEY'],
|
|
) as storage:
|
|
storage: S3Client
|
|
|
|
storage_bucket = config['STORAGE_BUCKET']
|
|
assert storage_bucket
|
|
|
|
# Make the bucket if it doesn't exist.
|
|
exists = await check_bucket_exists(storage, storage_bucket)
|
|
if not exists:
|
|
await storage.create_bucket(Bucket=storage_bucket)
|
|
print("Created bucket", storage_bucket)
|
|
else:
|
|
print("Bucket", storage_bucket, "already exists")
|
|
|
|
async with aiohttp.ClientSession() as session:
|
|
channels = [download_channel_attachments(session, storage, vision, channel_id) for channel_id in CHANNEL_IDS]
|
|
memes = await asyncio.gather(*channels)
|
|
|
|
|
|
asyncio.run(main())
|