feat: ✨ Add multi-threading to improve the process speed
This commit is contained in:
parent
23e7135004
commit
3deae416b2
90
main.py
90
main.py
@ -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)
|
||||
|
Reference in New Issue
Block a user