docker

main.py

@@ -4,18 +4,24 @@ import aiohttp
 import json
 import meilisearch
 
+from io import BytesIO
 from dataclasses import dataclass
-from typing import Dict, Any
+from typing import Dict, Any, List
 
 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
 
 # need to change the env for the google api calls
 load_dotenv()
 config = dotenv_values(".env")
 
+DATA_DIR = './data/memesearch'
+MEILI_INDEX = 'memes'
+MEILI_URL = 'http://localhost:7700'
+
 CHANNEL_IDS = [
     367936426962190336, # TPM shitposts
     1095060717117374687, # aella memes
@@ -31,6 +37,13 @@ headers = {
 # [ ] save img as json + imagefile
 # [ ] check extensions (ignore mp4 maybe)
 
+@dataclass
+class MemeDocument:
+    id: int
+    url: str
+    text: str | None
+    tags: List[str]
+
 @dataclass
 class MemeMessage:
     id: int
@@ -38,8 +51,8 @@ class MemeMessage:
     guild_name: str
     channel_name: str
 
-    text: str | None
     vision_json: str | None
+    text: str | None
 
     # maybe don't need these
     # kind of extra
@@ -47,22 +60,34 @@ class MemeMessage:
     guild_json: Dict[str, Any]
     channel_json: Dict[str, Any]
 
+    def output_dir(self):
+        return f'{DATA_DIR}/guilds/{self.guild_name}/{self.channel_name}'
+
+    def img_dir(self):
+        return f"{DATA_DIR}/imgs/"
+
     async def save_img(self, session: aiohttp.ClientSession):
-        output_dir = self.output_dir()
+        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(output_dir, exist_ok=True);
+
+        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:
-                with open(f'{output_dir}/{self.id}{ext}', 'wb') as f:
-                    f.write(await response.read())
+                res_bytes = await response.read()
 
-    def output_dir(self):
-        return f'./data/{self.guild_name}/{self.channel_name}'
+                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 = {
@@ -82,12 +107,13 @@ class MemeMessage:
         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 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(
@@ -142,9 +168,11 @@ async def download_channel_imgs(session, vision, channel_id):
     [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)
+    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)
 
 async def main():
     vision = google_vision.ImageAnnotatorClient()