implement db

README.md

@@ -0,0 +1,7 @@
+# Icon Generator for nKode
+
+## Install
+
+- download [wordnet](https://github.com/nltk/nltk_data/blob/gh-pages/packages/corpora/wordnet.zip)
+- `mkdir ~/nltk_data/corpora`
+- unzip wordnet and copy the folder into the new directory above

requirements.txt

@@ -1 +1,3 @@
 openai~=1.57.3
+requests~=2.32.3
+nltk~=3.9.1

src/db.py

@@ -1,44 +1,90 @@
-import sqlite3
+from sqlalchemy import create_engine, Column, String, Enum, func
-from pathlib import Path
+from sqlalchemy.orm import sessionmaker, declarative_base
 import uuid
+import hashlib
+from typing import Literal
+
+# Define the base for ORM mapping
+Base = declarative_base()
+
+# Define the Images table as a Python class
+class Image(Base):
+    __tablename__ = 'images'
+
+    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
+    prompt = Column(String, nullable=False)
+    model = Column(Enum("dall-e-3", "dall-e-2", name="model_enum"), nullable=False)
+    size = Column(String, nullable=False, default="1024x1024")
+    quality = Column(Enum("standard", "hd", name="quality_enum"), nullable=False, default="standard")
+    filename = Column(String, nullable=False)
+    hash = Column(String, nullable=True)
+
+# Create a database engine and session factory
+def create_db(db_path: str):
+    engine = create_engine(f'sqlite:///{db_path}')
+    Base.metadata.create_all(engine)
+    return engine
+
+def hash_image_row(
+        prompt: str,
+        model: str,
+        size: str,
+        quality: str,
+) -> str:
+    # Create a unique hash from the concatenated parameters
+    concat_string = f"{prompt}:{model}:{size}:{quality}"
+    unique_hash = hashlib.sha256(concat_string.encode()).hexdigest()
+    return unique_hash
 
 
-def create_db(db_path: Path):
+def insert_image_row(
-    conn = sqlite3.connect(db_path)
+        session,
-    c = conn.cursor()
+        image_id: str,
-
+        prompt: str,
-    # Create the table if it doesn't exist
+        hash: str,
-    c.execute("""
+        model: Literal["dall-e-3", "dall-e-2"] = "dall-e-3",
-    CREATE TABLE IF NOT EXISTS images (
+        size: Literal["1024x1024"] = "1024x1024",
-        id TEXT PRIMARY KEY,
+        quality: Literal["standard", "hd"] = "standard",
-        prompt TEXT,
+):
-        model TEXT,
+    image = Image(
-        size TEXT,
+        id=image_id,
-        quality TEXT,
+        prompt=prompt,
-        filename TEXT
+        model=model,
+        size=size,
+        quality=quality,
+        filename=f"{image_id}.png",
+        hash=hash
     )
-    """)
+    session.add(image)
-
+    session.commit()
-    # Commit and close the connection
-    conn.commit()
-    conn.close()
 
 
-def insert_into_db(db_path: Path, image_id: uuid, prompt: str, model: str = "dall-e-3", size: str = "1024x1024", quality: str = "standard"):
+def count_hash(session, hash_value: str):
-    conn = sqlite3.connect(db_path)
+    return session.query(func.count(Image.hash)).filter(Image.hash == hash_value).scalar()
-    c = conn.cursor()
+
-    # Insert the record into the database
+
-    c.execute(
+if __name__ == "__main__":
-        "INSERT INTO images (id, prompt, model, size, quality, filename) VALUES (?, ?, ?, ?, ?, ?)",
+    # Example usage
-        (
+    db_path = "../output/images.db"
-            image_id,
+    engine = create_db(db_path)
-            prompt,
+    Session = sessionmaker(bind=engine)
-            model,
+    session = Session()
-            size,
+
-            quality,
+    prompt = "A futuristic cityscape"
-            f"{image_id}.png"
+    model =  "dall-e-3"
-        )
+    size = "1024x1024"
+    quality = "standard"
+    hash = hash_image_row(prompt, model, size, quality)
+    # Insert an example record
+    insert_image_row(
+        session,
+        image_id=str(uuid.uuid4()),
+        prompt=prompt,
+        model=model,
+        size=size,
+        quality="standard",
+        hash=hash
     )
-    # Commit and close the connection
+
-    conn.commit()
+    session.close()
-    conn.close()
+

src/image_gen.py

@@ -45,3 +45,8 @@ def icon_gen(prompt: str, quality: Literal["hd", "standard"], output: Path = "./
     with open(image_filepath, "wb") as f:
         f.write(image_response.content)
     print(image_id)
+
+
+if __name__ == "__main__":
+    dog_emoji = image_style("dog in a hat on a beach with a drink", style=Styles.cartoon)
+    icon_gen(dog_emoji, output=Path("../output"), quality="hd")

src/wordlist.py

@@ -93,3 +93,18 @@ def is_imageable(noun):
                 return True
 
     return False
+
+
+if __name__ == "__main__":
+    nounlist = get_noun_list(Path("../data"))
+    banned_words = get_banned_wordlist(Path("../data"))
+    print(len(nounlist))
+    print(len( banned_words))
+    filtered_words = filter_banned_words_list(noun_list=nounlist, banned_words=banned_words)
+    print(len(filtered_words))
+
+    imageable = filter_nonimageable_words(filtered_words)
+    print(len(imageable))
+    print(f"imageable: {len(filtered_words) - len(imageable)}")
+    with open("../data/imageable.txt", "w") as fp:
+        fp.write("\n".join(imageable))

tests/image_gen_test.py

@@ -1,6 +0,0 @@
-from src.image_gen import image_style, Styles, icon_gen
-from pathlib import Path
-
-if __name__ == "__main__":
-    dog_emoji = image_style("dog in a hat on a beach with a drink", style=Styles.cartoon)
-    icon_gen(dog_emoji, output=Path("../output"), quality="hd")

tests/wordlist_test.py

@@ -1,16 +0,0 @@
-from src.wordlist import get_noun_list, get_banned_wordlist, filter_banned_words_list, filter_nonimageable_words
-from pathlib import Path
-
-if __name__ == "__main__":
-    nounlist = get_noun_list(Path("../data"))
-    banned_words = get_banned_wordlist(Path("../data"))
-    print(len(nounlist))
-    print(len( banned_words))
-    filtered_words = filter_banned_words_list(noun_list=nounlist, banned_words=banned_words)
-    print(len(filtered_words))
-
-    imageable = filter_nonimageable_words(filtered_words)
-    print(len(imageable))
-    print(f"imageable: {len(filtered_words) - len(imageable)}")
-    with open("../data/imageable.txt", "w") as fp:
-        fp.write("\n".join(imageable))