memestream bb

This commit is contained in:
Nathan Rashleigh committed 2024-04-08 00:34:30 +10:00
1 parent de58045c86
commit d82a7fdf1e
6 files changed
+145 -96

No files matched your search

+2 -1
View File
@@ -1,7 +1,8 @@
# MemeSearch # memestream
## TODO ## TODO
[ ] Backfill last N images from a channel [ ] Backfill last N images from a channel
[ ] Set oldest_img_id for channel [ ] Set oldest_img_id for channel
[ ] [ ]
[ ] change downloader to url type
+134 -77
View File
@@ -1,24 +1,35 @@
import os import os
import asyncio import asyncio
import aiohttp
import json import json
import meilisearch from urllib3.util import parse_url
from pathlib import Path
from io import BytesIO from io import BytesIO
from dataclasses import dataclass from dataclasses import dataclass
from uuid import uuid4, UUID
from typing import Dict, Any, List 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 dotenv import dotenv_values, load_dotenv
from google.cloud import vision as google_vision from google.cloud import vision as google_vision
from google.protobuf.json_format import MessageToDict, MessageToJson from google.protobuf.json_format import MessageToDict, MessageToJson
from google.cloud.vision_v1 import types 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 # need to change the env for the google api calls
load_dotenv() load_dotenv()
config = dotenv_values(".env") config = dotenv_values(".env")
DATA_DIR = './data/memesearch' DATA_DIR = './data/memestream'
MEILI_INDEX = 'memes' MEILI_INDEX = 'memes'
MEILI_URL = 'http://localhost:7700' MEILI_URL = 'http://localhost:7700'
@@ -39,57 +50,31 @@ headers = {
@dataclass @dataclass
class MemeDocument: class MemeDocument:
id: int id: UUID
url: str url: Url
text: str | None text: str
tags: List[str] tags: List[str]
@dataclass @dataclass
class MemeMessage: class DiscordAttachmentMessage:
id: int id: UUID
msg_id: int
url: str url: str
upload_url: str
guild_name: str guild_name: str
channel_name: str channel_name: str
vision_json: str | None vision_json: str | None
text: str | None
# maybe don't need these
# kind of extra
message_json: Dict[str, Any] message_json: Dict[str, Any]
guild_json: Dict[str, Any] guild_json: Dict[str, Any]
channel_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}' return f'{DATA_DIR}/guilds/{self.guild_name}/{self.channel_name}'
def img_dir(self): def file_dir(self):
return f"{DATA_DIR}/imgs/" return f"{DATA_DIR}/imgs/"
async def save_img(self, session: aiohttp.ClientSession): def detect_text(self, vision):
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 = { request = {
'image': { 'source': { 'image_uri': self.url } }, 'image': { 'source': { 'image_uri': self.url } },
'features': [types.Feature(type=types.Feature.Type.TEXT_DETECTION)], 'features': [types.Feature(type=types.Feature.Type.TEXT_DETECTION)],
@@ -104,24 +89,16 @@ class MemeMessage:
return request return request
def persist(self): 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) 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): def filter_message_img_urls(messages):
url_lists = map( 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) , messages)
return [item for group in url_lists for item in group] 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): async def fetch_channel_messages(channel_id):
url = f"https://discord.com/api/v9/channels/{channel_id}/messages" url = f"https://discord.com/api/v9/channels/{channel_id}/messages"
async with session.get(url, headers=headers) as res: 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: async with session.get(url, headers=headers) as res:
return await res.json() 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) channel = await fetch_channel(channel_id)
guild = await fetch_guild_details(channel['guild_id']) guild = await fetch_guild_details(channel['guild_id'])
print(f"Fetching messages for channel guild={guild['name']} name={channel['name']} id={channel['id']}") print(f"Fetching messages for channel guild={guild['name']} name={channel['name']} id={channel['id']}")
messages = await fetch_channel_messages(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( tasks = [download_attachment(session, storage, url, message) for (url, message) in urls_and_messages]
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) await asyncio.gather(*tasks)
[meme.annotate_image_request(vision) for meme in memes] # [attachment.detect_text(vision) for attachment in attachments]
[meme.persist() for meme in memes] # [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 # # have to vars the dataclass out for meili
documents = [vars(meme.to_document()) for meme in memes] # documents = [vars(attachment.to_meme()) for attachment in attachments]
client.index(MEILI_INDEX).add_documents(documents) # 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(): async def main():
assert config['MINIO_URL']
vision = google_vision.ImageAnnotatorClient() vision = google_vision.ImageAnnotatorClient()
async with aiohttp.ClientSession() as session:
channels = [download_channel_imgs(session, vision, channel_id) for channel_id in CHANNEL_IDS] # storage = Minio(config['MINIO_URL'],
memes = await asyncio.gather(*channels) # 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()) asyncio.run(main())
@@ -7,11 +7,11 @@
], ],
"build_systems": [ "build_systems": [
{ {
"name": "Python - memesearch", "name": "Python - memestream",
"cmd": ["$project_path/.venv/bin/python3", "-u", "main.py"] "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"] "cmd": ["$project_path/.venv/bin/python3", "-u", "$file"]
} }
] ]
+1
View File
@@ -4,3 +4,4 @@ google-cloud-vision
meilisearch meilisearch
pillow pillow
minio minio
aiobotocore
+1 -1
View File
@@ -15,7 +15,7 @@
const { InstantSearch, SearchBox, Hits, Highlight, Configure } = ReactInstantSearchDOM; const { InstantSearch, SearchBox, Hits, Highlight, Configure } = ReactInstantSearchDOM;
const searchClient = instantMeiliSearch( const searchClient = instantMeiliSearch(
"http://localhost:7700", "http://localhost:7700",
'73xjT1ueTlcXPEKJtzvJxtrlLNOR5q85fJ_TOzoRgcA' '33e714590199031d076954182c55a47b09ae7546e3dbf6f35862eb99cc950e1e'
); );
const App = () => ( const App = () => (
React.createElement(InstantSearch, { React.createElement(InstantSearch, {
+5 -15
View File
@@ -1,21 +1,16 @@
from minio import Minio from minio import Minio
from minio.error import S3Error
from dotenv import dotenv_values, load_dotenv def upload(config, document, source_file):
config = dotenv_values(".env") assert config['MINIO_URL']
def main():
client = Minio(config['MINIO_URL'], client = Minio(config['MINIO_URL'],
access_key=config['MINIO_ACCESS_KEY'], access_key=config['MINIO_ACCESS_KEY'],
secret_key=config['MINIO_SECRET_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 # The destination bucket and filename on the MinIO server
bucket_name = "memesearch" bucket_name = config['UPLOAD_BUCKET']
destination_file = "ok.png" destination_file = document.url
# Make the bucket if it doesn't exist. # Make the bucket if it doesn't exist.
found = client.bucket_exists(bucket_name) found = client.bucket_exists(bucket_name)
@@ -29,13 +24,8 @@ def main():
client.fput_object( client.fput_object(
bucket_name, destination_file, source_file, bucket_name, destination_file, source_file,
) )
print( print(
source_file, "successfully uploaded as object", source_file, "successfully uploaded as object",
destination_file, "to bucket", bucket_name, destination_file, "to bucket", bucket_name,
) )
if __name__ == "__main__":
try:
main()
except S3Error as exc:
print("error occurred.", exc)