Привет, Хабр!

Хочу поделиться своим опытом предобработки картиной с капчей и созданием модели, которая может определить, что же за символы в этой картинке. Код с архитектурой модели и обучением тоже будет, но основной упор часть с предобработкой картинок, поскольку это самая сложная часть. Также стоит упомянуть о допущениях, которые делались для упрощения задачи: использовались только латинские буквы (без цифр), все буквы в верхнем регистре, все капчи состоят из четырех символов (это самое серьезное допущение).

Капч в мире много, а собирать датасет парсером довольно проблематично. Да и размечать его потом очень не хочется. Поэтому мы будем использовать и для отложенной выборки, и для обучения картинки, сгенерированные библиотекой python captcha. Приложу пару примеров картинок, сгенерированных этой библиотекой:

Чтобы было проще воспроизвести код, прикладываю ссылку на requirements.txt на гите проекта: https://github.com/agusarev96/captcha_solver/blob/main/requirements.txt
Теперь перейдем для генерации датасета.
Импортируем библиотеки, которые нам нужны для того, чтобы сгенерировать нам большое количество картинок для обучения:

import random # поможет выбирать случайные символы
from captcha.image import ImageCaptcha # генерирует картинку
import string # отсюда берем список всех букв в кодировке ASCII

Запишем все наши символы, которые могут попасть в картинку в константу:

ALL_SYMBOLS = string.ascii_uppercase

При помощи библиотеки random сделаем 4 случайных выбора из всех наших символов и объединим это в строку:

def gen_captcha_text(length=4):

    symbols_list = []
    for _ in range(length):
        symbols_list.append(random.choice(ALL_SYMBOLS))

    return "".join(symbols_list)

Ну а дальше просто создаем 10000 картинок и сохраняем их в папку captcha_images. В названии файла прописываем текст, из которого генерировалась каптча (пригодится в дальнейшем):

# Create an image instance of the given size
image = ImageCaptcha(width = 280, height = 90)

for _ in range(10000):
    # Image captcha text
    captcha_text = gen_captcha_text(length=4)
     
    # generate the image of the given text
    data = image.generate(captcha_text)  

    # write the image on the given file and save it
    image.write(captcha_text, f'captcha_images/{captcha_text}.png')

Весь код с генерацией посмотреть можно тут: https://github.com/agusarev96/captcha_solver/blob/main/captcha_gen.py

"Можно приступать к обучению!", хотел бы сказать я, но не всё так просто. CNN решает нам задачу классификации (по одному символу). На выходе мы оставим len(ALL_SYMBOLS) нейронов, где каждому нейрону соответствует один символ. И решение модели будем принимать по самому высокому значению активации в нейроне. Следовательно, из каждой картинки нам надо вытащить область, соответствующую одному символу, что не так то просто. Причем делать нам это нужно и на этапе обучения, и на этапе применения модели. Перейдем к предобработке наших данных.

Создадим маппер, который мы будем использовать, чтобы складировать картинки в соответствующих папках:

target_mapper = {}
for index, value in enumerate(ALL_SYMBOLS):
    target_mapper[value] = index

Создадим две папки: training_data (обучающая выборка), test_data (тестовая). Тестовую выделяем для контроля переобучаемости модели.
Далее создадим скриптом по 26 папок (каждая соответствует одному символу) внутри training_data и test_data. В них мы и будем хранить картинки с вырезанными областями из картинки с капчей, соответствующими одному символу:

parent_dir = os.getcwd()

train_data = os.path.join(parent_dir, "training_data")
test_data = os.path.join(parent_dir, "test_data")

for symbol in target_mapper.keys():
    path = os.path.join(train_data, str(target_mapper[symbol]))
    try:
        os.mkdir(path)
    except:
        print(f"{path} already exists")

    path = os.path.join(test_data, str(target_mapper[symbol]))
    try:
        os.mkdir(path)
    except:
        print(f"{path} already exists")

Импортируем все, что нам потребуется для предобработки наших картинок:

from glob import glob
import cv2 as cv
import re
import numpy as np

glob позволит нам считать все пути до сгенерированных капч.

Теперь мы готовы к созданию нашей обучающей выборки: создадим список с путями до всех картинок. Далее создадим второй список такой же длины, где каждой выборке будет проставлен флаг тестовой или обучающей выборки (используем их в основном цикле, чтобы определять, в какую папку будем записывать результат).

all_captchas = glob("captcha_images/*.png")
sampleList = ["training_data", "test_data"]
 
randomList = random.choices(sampleList, weights=(80, 20), k=len(all_captchas))

Дальше я буду описывать основные части алгоритма, который применяется для каждой картинке, и буду описывать функции, которые использовались для предобработки.

Итак, первых шаг нашего алгоритма:
Считаем картинку при помощи библиотеки OpenCV (cv2) (достаточно просто передать путь до файла). Дальше переведем картинку в серые тона. Изначально для каждого пикселя в матрице лежит кортеж из трех величин от 0 до 255 с интенсивностью каждого цвета. Нам гораздо проще будет работать, если для каждого пикселя будет бинарная величина (окрашен или нет). Это сделает наши вычисления намного проще. Начнем с того, что преобразуем цветную картинку в серые тона (все ещё значение от 0 до 255, но одно на каждый пиксель).

img = cv.imread(all_captchas[idx])

# convert to grayscale
gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)

Дальше мы размажем немного картинку и создадим новый объект, где мы по-пиксельно разделим нашу картинку в серых тонах на заблюренную. Это поможет сразу в двух вещах: сделать единые контуры символов без разрывов и снизить уровень шума на картинках.

# blur
blur = cv.GaussianBlur(gray, (0, 0), sigmaX=33, sigmaY=33)
# divide
divide = cv.divide(gray, blur, scale=255)

Дальше, для простоты вычислений мы бинаризируем нашу картинку: вместо значений от 0 до 255 оставим только 0 и 255. Можно задавать границу вручную, но лучше работает алгоритм Отсу. Приложу пример из документации OpenCV: https://docs.opencv.org/4.x/otsu.jpg

# otsu threshold
thresh = cv.threshold(divide, 0, 255, cv.THRESH_BINARY+cv.THRESH_OTSU)[1] # 

Теперь проведем очистку шумов внутри наших символов, чтобы получить четкие контуры. Для того, чтобы сделать это, используются две операции: эрозия и расширение (Erosion и Dilation). Для этого создяется некоторое ядро, которое будет двигаться по картинке и производить одну из операций: для эрозии в ядре пиксель делается черным, если хоть один пиксель в ядре черный. Для расширения наоборот - если хотя бы один пиксель в ядре белый, то все ядро делается белым. Операция эрозии с последующим расширением хорошо помогает очистить шумы внешние. Расширение с последующей эрозией хорошо удаляет шумы внутри наших символов. Подробнее можно почитать и посмотреть тут: https://docs.opencv.org/4.x/d9/d61/tutorial_py_morphological_ops.html

# apply morphology
kernel = cv.getStructuringElement(cv.MORPH_RECT, (4, 4)) # создание ядра
morph = cv.morphologyEx(thresh, cv.MORPH_CLOSE, kernel) # dilation & erosion

Теперь проведем операцию, которая оставит нам только тонкую полоску на границе символа без заливки. Этот шаг не обязательный, но в моем случае помог оставить меньше внутренних контуров.

threshold = 150
canny_output = cv.Canny(morph, threshold, threshold * 2)
canny_output_cp = canny_output.copy()
canny_output = cv.blur(canny_output, (3, 3))

Теперь перейдем к той части, где можно сделать больше всего улучшений. Много в этой части сделано с целью получить на выходе ровно 4 области с символами. К сожалению, работает это не всегда четко, но без шума в данных для модели было бы даже не интересно.
Основные шаги, которые я применял, чтобы выдели эти 4 области:

1. Нашел контуры символов.

2. Для каждого контура нашел прямоугольную область, в которую он вписан.

3. Убрал контуры, которые полностью лежат внутри другого контура.

4. Отсортировал контуры по их площади.

5. Убрал контуры, которые больше 400 пикселей по площади (эта величина подбиралась для размера моих картинок).

6. Если у нас меньше четырех областей, я брал самый большой по площади, и делил его на две части. Нужно это для того, чтобы разъединить символы, которые наложились друг на друга и стали единым контуром.

7. Отсортировал наши области слева направа по координате левого верхнего угла.

8. Увеличил площадь каждого контура в полтора раза (алгоритм не мой). Этот шаг помогает в случаях, когда разделение пополам произошло не самым удачным образом. В картинку хоть и попадут части другого символа, зато будет целый символ.

9. Для первых четырех областей производим вырезание и сохраняем в нужную папку.

Итак, перейдем к коду. Для начала найдем контуры. Делается это при помощи встроенного метода в OpenCV:

contours, hierarchy = cv.findContours(
  canny_output, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE
)

Дальше найдем координаты верхнего левого и нижнего правого угла прямоугольника, который описывает каждый контур. Для этого для каждого контура пройдем по всем его точкам и найдем минимум и максимум для X и Y. Запишем это все в кортеж кортежей.

def get_corners(contour):
    min_x = contour[0][0][0]
    min_y = contour[0][0][1]
    max_x = contour[0][0][0]
    max_y = contour[0][0][1]
    for point in contour:
        x, y = point[0]
        min_x = min(min_x, x)
        min_y = min(min_y, y)
        max_x = max(max_x, x)
        max_y = max(max_y, y)

    return ((min_x, min_y), (max_x, max_y))

contours = [get_corners(i) for i in contours]

Теперь отфильтруем те боксы, которые полностью лежат внутри других. Область A включает в себя область B, если выполняются два условия:

1. Если координаты X и Y для левого верхнего угла первой области меньше либо равны этому же углу второй области.

2. Аналогично, правый нижний угол должен быть больше либо равен.

def inner_contour(boxB, boxA):
    startA, endA = boxA
    startB, endB = boxB
    start_lower = (startA[0] <= startB[0]) and (startA[1] <= startB[1])
    end_bigger = (endA[0] >= endB[0]) and (endA[1] >= endB[1])
    return start_lower and end_bigger

Произведем обход всех пар и уберем те, что являются для кого-то внутренним контуром.

def filter_inner_contour(boxes):
    result = []
    for i in range(len(boxes)):
        flag = False
        for j in range(len(boxes)):
            if j != i:
                if inner_contour(boxes[i], boxes[j]):
                    flag = True
        if not flag:
            result.append(boxes[i])
    return result

contours = filter_inner_contour(contours)

Сортировку произведем просто встроенной функцией python, передав ему функцию для расчета площади. Тут всё максимально просто. Здесь же уберем слишком большие области.

def get_square(start, end):
    length = end[0] - start[0]
    height = end[1] - start[1]
    return length * height

contours = sorted(contours, key=lambda x: get_square(*x), reverse=True)
contours = list(filter(lambda x: get_square(*x) > 400, contours))

Дальше произведем разбиение области на две части. Тут делаем сортировку по ширине нашей области. Делим самую большую на две другие области. O и Q на примере объединятся в один контур:

def split_rectangles(boxes):
    boxes = list(sorted(boxes, key=lambda x: x[1][0] - x[0][0]))
    box_to_split = boxes.pop()
    start, end = box_to_split
    width = end[0] - start[0]
    height = end[1] - start[1]
    box1 = ((start[0], start[1]), (start[0] + (width // 2), end[1]))
    box2 = ((start[0] + (width // 2), start[1]), (end[0], end[1]))
    boxes.append(box1)
    boxes.append(box2)
    return boxes

while len(contours) < 4:
    contours = split_rectangles(contours)

Теперь отсортируем по оси X для левого верхнего угла и произведем расширение области.

def expand_rectangle(rectangle, expansion_factor):
    start, end = rectangle
    x_min, y_min = start
    x_max, y_max = end
    width = x_max - x_min
    height = y_max - y_min
    x_center = (x_min + x_max) / 2
    y_center = (y_min + y_max) / 2

    new_width = width * expansion_factor
    new_height = height * expansion_factor

    x_min_expanded = max(int(x_center - new_width / 2), 0)
    y_min_expanded = max(int(y_center - new_height / 2), 0)
    x_max_expanded = int(x_center + new_width / 2)
    y_max_expanded = int(y_center + new_height / 2)

    return ((x_min_expanded, y_min_expanded), (x_max_expanded, y_max_expanded))

contours = sorted(contours, key=lambda x: x[0][0])
contours = [expand_rectangle(i, 1.5) for i in contours]

Вот так выглядит основное тело нашего цикла без описанных выше функций. Все вместе можно посмотреть в тетрадке на гите: https://github.com/agusarev96/captcha_solver/blob/main/captcha_solver_another_prepocessing.ipynb

for idx in range(len(all_captchas)):
    if randomList[idx]=="training_data":
        path = train_data
    else:
        path = test_data
        
    img = cv.imread(all_captchas[idx])

    symbols = re.search(r"\\(\w+)\.", all_captchas[idx]).group(1)
    for i in range(4):
        curr_path = os.path.join(path, str(target_mapper[symbols[i]]))
        curr_path = os.path.join(curr_path, "{0:010d}.png".format(curr_symb))
        curr_symb += 1
        crop_img = morph[contours[i][0][1]:contours[i][1][1], contours[i][0][0]:contours[i][1][0]]
        cv.imwrite(curr_path, crop_img)

Теперь перейдем к самой вкусной части - тренировки модели. Начнем со всех импортов:

import numpy as np 
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.metrics import roc_curve
from sklearn.metrics import roc_auc_score
from sklearn.metrics import auc
import keras
import tensorflow as tf
from keras import backend as K
from keras import metrics
from keras.regularizers import l2
from keras.models import Sequential
from keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout, BatchNormalization, InputLayer, Activation
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from keras.metrics import AUC
from keras.optimizers import Adam
from keras.callbacks import EarlyStopping
from sklearn.metrics import classification_report, confusion_matrix

Зафиксируем рандом.

seed_value= 145
np.random.seed(seed_value)
tf.random.set_seed(seed_value)
os.environ['PYTHONHASHSEED']=str(seed_value)

Пути с нашимим выборками.

train_path = os.path.join(os.getcwd(), "training_data")
test_path = os.path.join(os.getcwd(), "test_data")

Гиперпараметры модели.

hyper_dimension = 128
hyper_batch_size = 512
hyper_epochs = 100
hyper_channels = 1
hyper_mode = 'grayscale'

Создадим объекты генераторов картинок. Для трейна произведем небольшую аугментацию.

train_datagen = ImageDataGenerator(rescale=1.0/255.0, 
                                   shear_range = 0.2,
                                   zoom_range = 0.2, 
                                   horizontal_flip = True)
test_datagen = ImageDataGenerator(rescale=1.0/255.0) 

train_generator = train_datagen.flow_from_directory(directory = train_path, 
                                                    target_size = (hyper_dimension, hyper_dimension),
                                                    batch_size = hyper_batch_size, 
                                                    color_mode = hyper_mode,
                                                    class_mode = 'categorical',
                                                    # classes = target_mapper.values(),
                                                    seed = seed_value)
test_generator = test_datagen.flow_from_directory(directory = test_path, 
                                                 target_size = (hyper_dimension, hyper_dimension),
                                                 batch_size = hyper_batch_size, 
                                                 class_mode = 'categorical',
                                                  # classes = target_mapper.values(),
                                                 color_mode = hyper_mode,
                                                 shuffle=True,
                                                 seed = seed_value)

test_generator.reset()

Пропишем архитектуру нейронки.

model = Sequential()
model.add(InputLayer(input_shape=(hyper_dimension, hyper_dimension, hyper_channels)))

model.add(Conv2D(filters=32, kernel_size=3, activation="relu", kernel_regularizer=l2(0.05), bias_regularizer=l2(0.05)))
model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

model.add(Conv2D(filters=32, kernel_size=3, activation="relu", kernel_regularizer=l2(0.05), bias_regularizer=l2(0.05)))
model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

model.add(Conv2D(filters=32, kernel_size=3, activation="relu", kernel_regularizer=l2(0.05), bias_regularizer=l2(0.05)))
model.add(Flatten())

model.add(Dense(500, activation="relu", kernel_regularizer=l2(0.05), bias_regularizer=l2(0.05)))

# Output layer with 26 nodes (one for each possible letter/number we predict)
model.add(Dense(26, activation="softmax", kernel_regularizer=l2(0.05), bias_regularizer=l2(0.05)))

model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])

(Барабанная дробь). Ставим модель обучаться. Идем пить чай.

hist = model.fit(
    train_generator, 
    steps_per_epoch = len(train_generator),
    epochs = hyper_epochs,
    validation_data = test_generator,
    validation_steps = len(test_generator),
    verbose=2,
)

Не забываем сохранить модель.

model.save('cnn_captcha.h5')

Что же у нас получилось? Чтобы сделать инференс на новых данных нам нужно повторить те же процедуры с картинками, которые мы производили на этапе создания обучающего датасета. Тут я его повторять не буду. При желании, есть в гите. Стоит отметить подводный камень инференса нейронок: она будет возвращать N (N - количество классов. У нас 26) значений скора. Чем выше скор, тем выше её уверенность в том, что это соответствующий класс. НО!, классы она отсортировала не как числа, а как строки: первым будет класс 0, вторым 1, третьим - 10, а не 2. Учтите это на случай, если будете обучать свои модели (я сначала пошел переписывать весь алгоритм разделения на символы - решил, что на обучающей выборке выделялись нормально, но на новых данных всё пошло не так).
Итак, чтобы не было мешанины, создадим объект, который будет нам мапить обратно наши символы.

back_map = {}
for index, value in enumerate(ALL_SYMBOLS):
    back_map[str(index)] = value

classes = list(train_generator.class_indices.keys())

Передаем в цикле области с каждым символом, берем аргмакс по скору, передаем все в мапперы. Собираем для каждого символа решение модели.

predictions = m.predict(img_input, verbose=0)
predictions = predictions[0]

captcha_ls.append(back_map[classes[predictions.argmax()]])
    
result = "".join(captcha_ls)

Приведу фрагмент теста, который запускал для 50 картинок:

Что же в итоге? Получилось сделать модель, которая может впосле себе сностно предсказывать текст капчи. Да, это всего лишь MVP, который путает ряд символов между собой. Но это все равно неплохой результат, который предсказывает 4 символа с вероятностью ~70%.
Что можно улучшить?

1. Для начала можно добавить в капчу цифры.

2. Улучшить сам алгоритм нахождения символов.

3. Сделать алгоритм, который будет находить не фиксированное количество символов.

4. Сделать модель CRNN вместо CNN. Это как раз позволит нам не искать символы вручную и может предсказывать последовательности. (На этот счет постараюсь сделать ещё статью)