feat: Add multi-threading to improve the process speed

This commit is contained in:
Louis Gallet 2024-11-26 16:19:24 +01:00
parent 23e7135004
commit 3deae416b2
Signed by: lgallet
GPG Key ID: 84D3DF1528A84511

90
main.py
View File

@ -1,19 +1,45 @@
import string
from urllib.request import urlopen
from urllib.request import urlopen, urlretrieve
from PIL import Image, ImageDraw, ImageFont
from sys import argv
import random
import requests
import os
from multiprocessing import Pool
def getRandomFont(apiKey: str, usedFont: list=None) -> str:
CACHE_FOLDER = "cache/fonts"
# Créer le dossier de cache s'il n'existe pas
if not os.path.exists(CACHE_FOLDER):
os.makedirs(CACHE_FOLDER, exist_ok=True)
def get_font_path(font_url):
"""
Check if the font is already cached locally. If not, download it.
:param font_url: URL of the TTF font file.
:return: Path to the locally cached TTF file.
"""
font_name = font_url.split("/")[-1] # Extract the font filename from the URL
local_path = os.path.join(CACHE_FOLDER, font_name)
if not os.path.exists(local_path):
print(f"Downloading font: {font_name}")
urlretrieve(font_url, local_path) # Download and save the file locally
else:
print(f"Using cached font: {font_name}")
return local_path
def getRandomFont(apiKey: str, usedFont: list = None) -> str:
"""
Function to fetch a random font from a list of available fonts
provided by Google Font
:param apiKey: API Key used to connect to Google Fonts.
:param usedFont: List of already used fonts. Set to none by default
:return: Link to the ttf file
:param usedFont: List of already used fonts. Set to None by default
:return: Local path to the cached TTF file.
"""
if usedFont is None:
usedFont = []
@ -26,7 +52,7 @@ def getRandomFont(apiKey: str, usedFont: list=None) -> str:
if "items" not in fonts_data:
raise Exception("No fonts found in the API response.")
available_fonts = fonts_data["items"]
unused_fonts = [font for font in available_fonts if usedFont is None or font["family"] not in usedFont]
unused_fonts = [font for font in available_fonts if font["family"] not in usedFont]
if not unused_fonts:
raise Exception("No unused fonts available to select from.")
@ -41,15 +67,15 @@ def getRandomFont(apiKey: str, usedFont: list=None) -> str:
raise Exception(f"No TTF link found for the font: {chosen_font['family']}")
print(f"Selected font: {chosen_font['family']}")
return ttf_link
return get_font_path(ttf_link)
def create_letter_image(letter, output_path, font):
def create_letter_image(letter, output_path, font_path):
"""
Function to create an image for a letter passed as an argument
:param letter: the letter
:param output_path: the path to save the image
:param font: font to be used (should be an online Font URL)
:param font_path: local path to the TTF font file
:return: print a message and save the image
"""
width, height = 13, 13
@ -59,9 +85,9 @@ def create_letter_image(letter, output_path, font):
draw = ImageDraw.Draw(image)
try:
font = ImageFont.truetype(urlopen(font), font_size)
font = ImageFont.truetype(font_path, font_size)
except IOError:
raise Exception("Cannot open font URL")
raise Exception(f"Cannot load font file: {font_path}")
bbox = draw.textbbox((0, 0), letter, font=font)
text_width, text_height = bbox[2] - bbox[0], bbox[3] - bbox[1]
@ -74,29 +100,41 @@ def create_letter_image(letter, output_path, font):
image.save(output_path)
print(f"Image saved at {output_path}")
def createImageForEachLetter(output_folder, font, index=0):
def createImageForEachLetter(args):
"""
Function to create an image for each of the letters passed as an argument
:param output_folder: the folder where to save the image
:param font: Google Font URL
:param index: the current index of the iteration (default to 0 if there is no iteration)
:return: Save the image into the correct subfolder
:param args: Tuple containing output folder, font path, and index
"""
alphabet = string.ascii_letters
if not os.path.exists(output_folder):
os.makedirs(output_folder)
for element in alphabet:
if not os.path.exists(os.path.join(output_folder, element)):
os.makedirs(os.path.join(output_folder, element))
create_letter_image(element, f"{output_folder}/{element}/{element}-{index}.png", font)
output_folder, font_path, index = args
alphabet = string.ascii_letters # Inclut à la fois les lettres minuscules et majuscules
os.makedirs(output_folder, exist_ok=True) # Ensure base output folder exists
for element in alphabet:
letter_folder = os.path.join(output_folder, element)
os.makedirs(letter_folder, exist_ok=True)
create_letter_image(element, f"{letter_folder}/{element}-{index}.png", font_path)
if __name__ == "__main__":
if len(argv) < 3:
raise Exception("Usage: " + argv[0] + " <number-of-letter> <api-key>")
raise Exception("Usage: " + argv[0] + " <number-of-iterations> <api-key>")
if input(f"This script will generate {argv[1]} images for each letter in the alphabet. Continue? Y/n ") != "Y":
print("Exiting...")
exit(0)
# Configuration
num_iterations = int(argv[1])
api_key = argv[2]
usedFont = []
for i in range(int(argv[1])):
createImageForEachLetter("dataset", getRandomFont(argv[2], usedFont), i)
# Fetch fonts and cache them
fonts = [getRandomFont(api_key, usedFont) for _ in range(num_iterations)]
# Prepare arguments for multiprocessing
args = [("dataset", fonts[i], i) for i in range(num_iterations)]
# Use multiprocessing pool
with Pool() as pool:
pool.map(createImageForEachLetter, args)