Refactor; move consts to consts.py

.gitignore

@@ -1,5 +1,7 @@
-font_images
+generated_images
 train_test_images
 .DS_Store
 .ipynb_checkpoints/visualize-checkpoint.ipynb
 font_identifier_model.pth
+*.pyc
+__pycache__

arrange_train_test_images.py

@@ -3,18 +3,17 @@
 import os
 import shutil
 import random
+from consts import TRAIN_TEST_IMAGES_DIR, GEN_IMAGES_DIR
 
-source_dir = './font_images'
-organized_dir = './train_test_images'
-train_dir = os.path.join(organized_dir, 'train')
-test_dir = os.path.join(organized_dir, 'test')
+train_dir = os.path.join(TRAIN_TEST_IMAGES_DIR, 'train')
+test_dir = os.path.join(TRAIN_TEST_IMAGES_DIR, 'test')
 
 # create directories if they don't exist
 os.makedirs(train_dir, exist_ok=True)
 os.makedirs(test_dir, exist_ok=True)
 
 # make a list of all the font names
-fonts = [f.split('_')[0] for f in os.listdir(source_dir) if f.endswith('.png')]
+fonts = [f.split('_')[0] for f in os.listdir(GEN_IMAGES_DIR) if f.endswith('.png')]
 fonts = list(set(fonts))  # getting unique font names
 
 for font in fonts:
@@ -23,7 +22,7 @@ for font in fonts:
     os.makedirs(font_train_dir, exist_ok=True)
     os.makedirs(font_test_dir, exist_ok=True)
 
-    font_files = [f for f in os.listdir(source_dir) if f.startswith(font)]
+    font_files = [f for f in os.listdir(GEN_IMAGES_DIR) if f.startswith(font)]
     random.shuffle(font_files)
 
     train_files = font_files[:int(0.8 * len(font_files))]
@@ -31,8 +30,8 @@ for font in fonts:
 
     # Move training files
     for train_file in train_files:
-        shutil.move(os.path.join(source_dir, train_file), font_train_dir)
+        shutil.move(os.path.join(GEN_IMAGES_DIR, train_file), font_train_dir)
 
     # Move test files
     for test_file in test_files:
-        shutil.move(os.path.join(source_dir, test_file), font_test_dir)
+        shutil.move(os.path.join(GEN_IMAGES_DIR, test_file), font_test_dir)

consts.py

@@ -0,0 +1,10 @@
+# number of images to generate per font
+IMAGES_PER_FONT = 50
+# allowlist of fonts to use
+FONT_ALLOWLIST = ["Arial", "Avenir", "Courier", "Helvetica", "Georgia", "Tahoma", "Times New Roman", "Verdana"]
+# directory where to store the generated images
+GEN_IMAGES_DIR = './generated_images'
+# images organized into train and test directories
+TRAIN_TEST_IMAGES_DIR = './train_test_images'
+# where to grab the font files from
+FONT_FILE_DIRS = ['/System/Library/Fonts/', '/System/Library/Fonts/Supplemental/']

gen_sample_data.py

@@ -6,22 +6,15 @@ from PIL import Image, ImageDraw, ImageFont
 import nltk
 from nltk.corpus import brown
 import random
-
-IMAGES_PER_FONT = 50
+from consts import FONT_ALLOWLIST, IMAGES_PER_FONT, GEN_IMAGES_DIR, FONT_FILE_DIRS
 
 # Download the necessary data from nltk
 nltk.download('brown')
 
-# Note that this will only work on MacOS where this is the default font directory
-font_dirs = ['/System/Library/Fonts/', '/System/Library/Fonts/Supplemental/']
-output_dir = './font_images'
-os.makedirs(output_dir, exist_ok=True)
+os.makedirs(GEN_IMAGES_DIR, exist_ok=True)
 
 all_brown_words = sorted(set(brown.words(categories='news')))
 
-# This is a list of fonts that we want to use for our sample data
-FONT_ALLOWLIST = ["Arial", "Avenir", "Courier", "Helvetica", "Georgia", "Tahoma", "Times New Roman", "Verdana"]
-
 def wrap_text(text, line_length=10):
     """Wraps the provided text every 'line_length' words."""
     words = text.split()
@@ -37,38 +30,42 @@ def random_code_text(base_code, num_lines=15):
     lines = base_code.split("\n")
     return "\n".join(random.sample(lines, min(num_lines, len(lines))))
 
-for font_dir in font_dirs:
-    for font_file in os.listdir(font_dir):
-        if font_file.endswith('.ttf') or font_file.endswith('.ttc'):
-            font_path = os.path.join(font_dir, font_file)
-            font_name = font_file.split('.')[0]
-            if font_name not in FONT_ALLOWLIST:
-                continue
-            # Output the font name so we can see the progress
-            print(font_path, font_name)
+def main():
+    for font_dir in FONT_FILE_DIRS:
+        for font_file in os.listdir(font_dir):
+            if font_file.endswith('.ttf') or font_file.endswith('.ttc'):
+                font_path = os.path.join(font_dir, font_file)
+                font_name = font_file.split('.')[0]
+                if font_name not in FONT_ALLOWLIST:
+                    continue
+                # Output the font name so we can see the progress
+                print(font_path, font_name)
+
+                if font_file.endswith('.ttc'):
+                    # ttc fonts have multiple fonts in one file, so we need to specify which one we want
+                    font = ImageFont.truetype(font_path, random.choice(range(32, 128)), index=0)
+                else:
+                    # ttf fonts have only one font in the file
+                    font_size = random.choice(range(32, 128))  # Increased minimum font size
+                    font = ImageFont.truetype(font_path, font_size)
 
-            if font_file.endswith('.ttc'):
-                # ttc fonts have multiple fonts in one file, so we need to specify which one we want
-                font = ImageFont.truetype(font_path, random.choice(range(32, 128)), index=0)
-            else:
-                # ttf fonts have only one font in the file
-                font_size = random.choice(range(32, 128))  # Increased minimum font size
-                font = ImageFont.truetype(font_path, font_size)
+                # Counter for the image filename
+                j = 0
+                for i in range(IMAGES_PER_FONT):  # Generate 50 images per font - reduced to 10 for now to make things faster
+                    prose_sample = random_prose_text(all_brown_words)
 
-            # Counter for the image filename
-            j = 0
-            for i in range(IMAGES_PER_FONT):  # Generate 50 images per font - reduced to 10 for now to make things faster
-                prose_sample = random_prose_text(all_brown_words)
+                    for text in [prose_sample]:
+                        img = Image.new('RGB', (800, 400), color="white")  # Canvas size
+                        draw = ImageDraw.Draw(img)
 
-                for text in [prose_sample]:
-                    img = Image.new('RGB', (800, 400), color="white")  # Canvas size
-                    draw = ImageDraw.Draw(img)
+                        # Random offsets, but ensuring that text isn't too far off the canvas
+                        offset_x = random.randint(-20, 10)
+                        offset_y = random.randint(-20, 10)
+                        draw.text((offset_x, offset_y), text, fill="black", font=font)
 
-                    # Random offsets, but ensuring that text isn't too far off the canvas
-                    offset_x = random.randint(-20, 10)
-                    offset_y = random.randint(-20, 10)
-                    draw.text((offset_x, offset_y), text, fill="black", font=font)
+                        j += 1
+                        output_file = os.path.join(GEN_IMAGES_DIR, f"{font_name}_{j}.png")
+                        img.save(output_file)
 
-                    j += 1
-                    output_file = os.path.join(output_dir, f"{font_name}_{j}.png")
-                    img.save(output_file)
+if __name__ == '__main__':
+    main()

train_font_identifier.py

@@ -1,15 +1,11 @@
-import copy
 import os
-import time
 import torch
 import torch.optim as optim
 import torch.nn as nn
-from torch.optim import lr_scheduler
 from torchvision import datasets, models, transforms
 from tqdm import tqdm
-
-# Directory with organized font images
-data_dir = './train_test_images'
+import torch
+from consts import TRAIN_TEST_IMAGES_DIR
 
 # Transformations for the image data
 data_transforms = transforms.Compose([
@@ -21,7 +17,7 @@ data_transforms = transforms.Compose([
 
 # Create datasets
 image_datasets = {
-    x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms)
+    x: datasets.ImageFolder(os.path.join(TRAIN_TEST_IMAGES_DIR, x), data_transforms)
     for x in ['train', 'test']
 }
 
@@ -92,3 +88,5 @@ for epoch in range(num_epochs):
 
 # Save the model to disk
 torch.save(model.state_dict(), 'font_identifier_model.pth')
+
+