diff --git a/TODO.md b/TODO.md index 5d7f3de..2820fb4 100644 --- a/TODO.md +++ b/TODO.md @@ -1,7 +1,8 @@ -# MemeSearch +# memestream ## TODO [ ] Backfill last N images from a channel [ ] Set oldest_img_id for channel [ ] +[ ] change downloader to url type diff --git a/main.py b/main.py index 26ccfc1..64cab5f 100644 --- a/main.py +++ b/main.py @@ -1,24 +1,35 @@ import os import asyncio -import aiohttp import json -import meilisearch - +from urllib3.util import parse_url +from pathlib import Path from io import BytesIO from dataclasses import dataclass +from uuid import uuid4, UUID from typing import Dict, Any, List +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 -from PIL import Image +from minio import Minio +import meilisearch +from urllib3.util import Url, url + +import uploader # need to change the env for the google api calls load_dotenv() config = dotenv_values(".env") -DATA_DIR = './data/memesearch' +DATA_DIR = './data/memestream' MEILI_INDEX = 'memes' MEILI_URL = 'http://localhost:7700' @@ -39,57 +50,31 @@ headers = { @dataclass class MemeDocument: - id: int - url: str - text: str | None + id: UUID + url: Url + text: str tags: List[str] @dataclass -class MemeMessage: - id: int +class DiscordAttachmentMessage: + id: UUID + msg_id: int url: str + upload_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): + def document_dir(self): return f'{DATA_DIR}/guilds/{self.guild_name}/{self.channel_name}' - def img_dir(self): + def file_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): + def detect_text(self, vision): request = { 'image': { 'source': { 'image_uri': self.url } }, 'features': [types.Feature(type=types.Feature.Type.TEXT_DETECTION)], @@ -104,24 +89,16 @@ class MemeMessage: return request def persist(self): - with open(f'{self.output_dir()}/{self.id}.json', 'w') as f: + with open(f'{self.document_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'])) + 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_imgs(session, vision, channel_id): +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: @@ -142,42 +119,122 @@ async def download_channel_imgs(session, vision, channel_id): async with session.get(url, headers=headers) as res: return await res.json() + async def download_attachment( + session: aiohttp.ClientSession, + storage: S3Client, + url: Url, + message, + ): + # there's a big param of some kind after the image url that we remove + if not url.path: + raise RuntimeError(f"url missing path url={url}") + + storage_bucket = config['UPLOAD_BUCKET'] + if not storage_bucket: + raise RuntimeError(f"config missing UPLOAD_BUCKET") + + (filename, ext) = os.path.splitext(os.path.basename(url.path)) + # TODO + # FILTER WANTED TYPES + + id = uuid4() + + async with session.get(url.url) as response: + if response.status == 200: + + # make a webp conversion for obvious image types + convert_extensions = ['jpg', 'png', 'jpeg', 'webp'] + res_bytes = await response.read() + object_name: str + + # trim the . + if ext.lower()[1:] in convert_extensions: + image = Image.open(BytesIO(res_bytes)) + + webp_output = BytesIO() + image.save(webp_output, format='WebP') + webp_output.seek(0) + + object_name = f"{id}.webp" + print(f"uploading object_name={object_name} filename={filename} ext={ext}") + return await storage.put_object(Bucket=storage_bucket, Key=object_name, Body=webp_output) + + object_name = f"{id}{ext}" + print(f"uploading object_name={object_name} filename={filename} ext={ext}") + return await storage.put_object(Bucket=storage_bucket, Key=object_name, Body=res_bytes) + + # add discord id to server processed msgs list + # if there then nothing to do + + # return MemeDocument( + # id=id, + # url=self.upload_url, + # text=self.text or '', + # tags=[self.guild_name, self.channel_name], + # ) + 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) + urls_and_messages = 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] + tasks = [download_attachment(session, storage, url, message) for (url, message) in urls_and_messages] await asyncio.gather(*tasks) - [meme.annotate_image_request(vision) for meme in memes] - [meme.persist() for meme in memes] + # [attachment.detect_text(vision) for attachment in attachments] + # [attachment.persist() for attachment in attachments] - client = meilisearch.Client(MEILI_URL, config['MEILI_MASTER_KEY']) + # 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) + # # 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['MINIO_URL'] + 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) + + # storage = Minio(config['MINIO_URL'], + # access_key=config['MINIO_ACCESS_KEY'], + # secret_key=config['MINIO_SECRET_KEY'], + # ) + + 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['UPLOAD_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()) diff --git a/memesearch.sublime-project b/memestream.sublime-project similarity index 73% rename from memesearch.sublime-project rename to memestream.sublime-project index a4f79a1..c5a0ad2 100644 --- a/memesearch.sublime-project +++ b/memestream.sublime-project @@ -7,11 +7,11 @@ ], "build_systems": [ { - "name": "Python - memesearch", + "name": "Python - memestream", "cmd": ["$project_path/.venv/bin/python3", "-u", "main.py"] }, { - "name": "Python - memesearch - current file", + "name": "Python - memestream - current file", "cmd": ["$project_path/.venv/bin/python3", "-u", "$file"] } ] diff --git a/requirements.txt b/requirements.txt index ef5dce7..fd0437a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,4 +3,5 @@ aiohttp google-cloud-vision meilisearch pillow -minio \ No newline at end of file +minio +aiobotocore \ No newline at end of file diff --git a/search.html b/search.html index 0a691bb..4ef7cec 100644 --- a/search.html +++ b/search.html @@ -15,7 +15,7 @@ const { InstantSearch, SearchBox, Hits, Highlight, Configure } = ReactInstantSearchDOM; const searchClient = instantMeiliSearch( "http://localhost:7700", - '73xjT1ueTlcXPEKJtzvJxtrlLNOR5q85fJ_TOzoRgcA' + '33e714590199031d076954182c55a47b09ae7546e3dbf6f35862eb99cc950e1e' ); const App = () => ( React.createElement(InstantSearch, { diff --git a/file_uploader.py b/uploader.py similarity index 65% rename from file_uploader.py rename to uploader.py index b08de53..faac011 100644 --- a/file_uploader.py +++ b/uploader.py @@ -1,21 +1,16 @@ from minio import Minio -from minio.error import S3Error -from dotenv import dotenv_values, load_dotenv -config = dotenv_values(".env") +def upload(config, document, source_file): + assert config['MINIO_URL'] -def main(): client = Minio(config['MINIO_URL'], access_key=config['MINIO_ACCESS_KEY'], secret_key=config['MINIO_SECRET_KEY'], ) - # The file to upload, change this path if needed - source_file = "./9a6.png" - # The destination bucket and filename on the MinIO server - bucket_name = "memesearch" - destination_file = "ok.png" + bucket_name = config['UPLOAD_BUCKET'] + destination_file = document.url # Make the bucket if it doesn't exist. found = client.bucket_exists(bucket_name) @@ -29,13 +24,8 @@ def main(): client.fput_object( bucket_name, destination_file, source_file, ) + print( source_file, "successfully uploaded as object", destination_file, "to bucket", bucket_name, - ) - -if __name__ == "__main__": - try: - main() - except S3Error as exc: - print("error occurred.", exc) \ No newline at end of file + ) \ No newline at end of file