feat: ✨ Create the main program
This commit is contained in:
parent
864a442e9e
commit
3f97bb69c1
1
.gitignore
vendored
1
.gitignore
vendored
@ -12,6 +12,7 @@
|
||||
.idea/**/dictionaries
|
||||
.idea/**/shelf
|
||||
|
||||
.idea
|
||||
# AWS User-specific
|
||||
.idea/**/aws.xml
|
||||
|
||||
|
3
README.md
Normal file
3
README.md
Normal file
@ -0,0 +1,3 @@
|
||||
# OCR Dataset Generator
|
||||
Python script that generates datasets for the recognition of text written in different fonts.
|
||||
The images generated are in 13x13 format and are all arranged in alphabetical order.
|
102
main.py
Normal file
102
main.py
Normal file
@ -0,0 +1,102 @@
|
||||
import string
|
||||
from urllib.request import urlopen
|
||||
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
from sys import argv
|
||||
import random
|
||||
import requests
|
||||
import os
|
||||
|
||||
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
|
||||
"""
|
||||
if usedFont is None:
|
||||
usedFont = []
|
||||
url = f"https://www.googleapis.com/webfonts/v1/webfonts?key={apiKey}"
|
||||
response = requests.get(url)
|
||||
if response.status_code != 200:
|
||||
raise Exception(f"Failed to fetch fonts: {response.status_code} - {response.reason}")
|
||||
|
||||
fonts_data = response.json()
|
||||
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]
|
||||
|
||||
if not unused_fonts:
|
||||
raise Exception("No unused fonts available to select from.")
|
||||
chosen_font = random.choice(unused_fonts)
|
||||
|
||||
usedFont.append(chosen_font["family"])
|
||||
|
||||
font_files = chosen_font.get("files", {})
|
||||
ttf_link = font_files.get("regular") or next(iter(font_files.values()), None)
|
||||
|
||||
if not ttf_link:
|
||||
raise Exception(f"No TTF link found for the font: {chosen_font['family']}")
|
||||
|
||||
print(f"Selected font: {chosen_font['family']}")
|
||||
return ttf_link
|
||||
|
||||
|
||||
def create_letter_image(letter, output_path, font):
|
||||
"""
|
||||
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)
|
||||
:return: print a message and save the image
|
||||
"""
|
||||
width, height = 13, 13
|
||||
font_size = 12
|
||||
|
||||
image = Image.new('L', (width, height), color=255) # Fond blanc
|
||||
draw = ImageDraw.Draw(image)
|
||||
|
||||
try:
|
||||
font = ImageFont.truetype(urlopen(font), font_size)
|
||||
except IOError:
|
||||
raise Exception("Cannot open font URL")
|
||||
|
||||
bbox = draw.textbbox((0, 0), letter, font=font)
|
||||
text_width, text_height = bbox[2] - bbox[0], bbox[3] - bbox[1]
|
||||
|
||||
text_x = (width - text_width) // 2
|
||||
text_y = (height - text_height) // 2
|
||||
|
||||
draw.text((text_x, text_y), letter, fill=0, font=font) # Texte noir
|
||||
|
||||
image.save(output_path)
|
||||
print(f"Image saved at {output_path}")
|
||||
|
||||
def createImageForEachLetter(output_folder, font, index=0):
|
||||
"""
|
||||
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
|
||||
"""
|
||||
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)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if len(argv) < 3:
|
||||
raise Exception("Usage: " + argv[0] + " <number-of-letter> <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)
|
||||
usedFont = []
|
||||
for i in range(int(argv[1])):
|
||||
createImageForEachLetter("dataset", getRandomFont(argv[2], usedFont), i)
|
Reference in New Issue
Block a user