Всем привет, это моя первая статья на Хабре. В этой статье я хочу рассказать, как сгенерировать датасет печатных букв с помощью .ttf файла и кода на Python в 170 строк.
Зачем?
Для начала выясним, зачем нужно генерировать датасет. В моем случае стоял пользовательский интерес в распознавании шифров с конструкторских документов при сканировании чертежей. Для распознавания нужна нейросеть, чтобы обучить нейросеть нужен датасет. В общем, нейросеть без датасета как машина без колес, а значит необходимо иметь возможность генерировать датасет. При попытке найти датасет по типу EMNIST было обнаружено, что печатные шрифты никто не выкладывает, по крайней мере, на Кириллице. А если нужно больше данных чем есть, или нужно поменять искажения текста. Так и было принято решение написать генератор датасета по типу EMNIST. Но с печатными буквами русского алфавита.
Создание изображения
Начнем с библиотек, нам потребуется следующий инструментарий:
from PIL import Image, ImageDraw, ImageFont, ImageFilter import os from glob import glob import numpy as npPIL — он же Pillow, позволит нам создавать изображения и работать с ними.
OS — позволит работать с директориями.
glob — для поиска изображений в папках.
numpy — для работы с массивами.
Я работаю в PyCharm, поэтому не буду затрагивать тему установки библиотек. Об этом можно почитать на официальном сайте Python.
Что мы имеем на входе?
На входе нам даны следующие параметры:
FONT_PATH = "GOST2304A.ttf" # Путь к файлу шрифта (поддерживающему кириллицу) FONT_SIZE = 36 # Размер шрифта IMAGE_SIZE = (28, 28) # Размер выходного изображения BACKGROUND_COLOR = (255, 255, 255) # Белый фон TEXT_COLOR = (0, 0, 0) # Черный цвет текста NUM_IMAGES = 10 # Количество изображений на букву # Список русских букв RUSSIAN_LETTERS = [ 'А', 'Б', 'В', 'Г', 'Д', 'Е', 'Ё', 'Ж', 'З', 'И', 'Й', 'К', 'Л', 'М', 'Н', 'О', 'П', 'Р', 'С', 'Т', 'У', 'Ф', 'Х', 'Ц', 'Ч', 'Ш', 'Щ', 'Ъ', 'Ы', 'Ь', 'Э', 'Ю', 'Я' ]Как видно из кода, мы можем регулировать следующие параметры:
-
шрифт
-
размер шрифта
-
размер картинки (в пикселях)
-
цвет фона (в RGB)
-
цвет текста (в RGB)
-
количество экземпляров каждой буквы в нашем будущем датасете.
-
Алфавит для генерации
3. Дальше напишем функцию для генерации картинок, на вход она будет принимать только букву и путь до шрифта:
def create_letter_image(letter, font_path): """Создает изображение с наклонной буквой""" # Загружаем шрифт try: base_font = ImageFont.truetype(font_path, FONT_SIZE) except OSError as e: print(f"Ошибка загрузки шрифта: {font_path}") print(e) return # Создаем временное изображение для измерения текста temp_image = Image.new('RGB', IMAGE_SIZE, BACKGROUND_COLOR) draw = ImageDraw.Draw(temp_image) text_width, text_height = draw.textbbox((0, 0), letter, font=base_font)[2:] # Создаем основное изображение с учетом наклона image = Image.new('RGB', (int(IMAGE_SIZE[0]), IMAGE_SIZE[1]), BACKGROUND_COLOR) draw = ImageDraw.Draw(image) # Рисуем текст с наклоном x = (image.width - text_width) / 2 y = (image.height - text_height) / 2 draw.text((x, y), letter, font=base_font, fill=TEXT_COLOR) return imageСначала проверяем наличие шрифта в нашем проекте.
Создаем новое изображение с помощью Image.new. Задаем цветовой канал, размер изображения и цвет фона.
С помощью метода textbbox создадим поле для написания текста и зададим высоту и ширину поля.
Задаем координаты центра буквы и заполняем текстовое поле взяв букву из нашего алфавита.
Добавляем шумы
Для создания уникальных изображений и во избежание повторов при генерации и как следствии получении не репрезентативной выборки создадим функцию, которая будет генерировать шум на изображении случайно добавляя пиксели разной яркости (и здесь мы будем конвертировать нашего изображение из RGB в оттенки серого).
def add_noise(image): """ Добавляет гауссов шум на изображение в оттенках серого (шум одинаков для всех каналов, сохраняя изображение серым) """ # Конвертируем в numpy array и преобразуем в grayscale img_array = np.array(image.convert('L')) # 'L' - режим оттенков серого height, width = img_array.shape noisy = img_array.copy() # Параметры гауссова шума mean = 0 var = np.random.uniform(0.001, 0.02) # Диапазон дисперсии sigma = var ** 0.5 # Генерируем шум (один канал) gauss = np.random.normal(mean, sigma, (height, width)) # Применяем шум и обрезаем значения noisy = np.clip(noisy + gauss * 255, 0, 255).astype(np.uint8) # Конвертируем обратно в RGB (но сохраняем оттенки серого) return Image.fromarray(noisy).convert('RGB')Было использовано гауссовское распределение для отрисовки серых пикселей случайной яркости.
Итоговая функция
Итоговая функция будет выглядеть следующим образом:
def main(): # Загружаем шрифт try: font = ImageFont.truetype(FONT_PATH, FONT_SIZE) except IOError: print(f"Ошибка: Шрифт по пути '{FONT_PATH}' не найден") return # Создаем корневую директорию for letter in RUSSIAN_LETTERS: # Создаем директорию для буквы letter_dir = f"Generated_images/Letter_{letter}" os.makedirs(letter_dir, exist_ok=True) # Генерируем эталонное изображение # Сохраняем NUM_IMAGES копий for i in range(NUM_IMAGES): img = create_letter_image(letter, FONT_PATH) img = add_noise(img) file_name = f"Letter_{letter}_{i:04d}.png" file_path = os.path.join(letter_dir, file_name) img.save(file_path) print(f"Сгенерировано {NUM_IMAGES} изображений для буквы '{letter}'")После написания всех функций осталось только пробежаться по всем буквам нашей кириллицы и создать соответствующее количество изображений, после чего в нашем проекте появиться соответствующая папка с нашим датасетом, где имя каждой папки будет соответствовать будущему классу буквы.
Итог
Надеюсь этот пример будет полезен тем, кто хочет обучать свои небольшие нейросети для узких задач, но кто упирается в отсутствие нужных датасетов. При желании можно добавить функцию деформации изображения для улучшения распознавания, например, наклонных отсканированных документов, но это уже будет описано в следующей статье.
Подробный код можно скачать с моего GitHub.
ссылка на оригинал статьи https://habr.com/ru/articles/925022/
Добавить комментарий