Однажды я решил проверить, можно ли решать судоку в реальном времени с помощью машинного обучения и обработки изображений. Я изначально думал, что это можно сделать за 2-3 часа, просто потому, что готовые библиотеки можно найти под любые нужды. Спойлер: я ужасно ошибся в своей оценке, но в итоге это стало интересным уроком.
Судоку — это логическая игра. Идея игры состоит в том, чтобы размещать числа на сетке 9x9 так, чтобы каждый столбец, строка и каждая из подсеток 3x3 содержали все цифры от 1 до 9. Поиск решения судоку в режиме реального времени потребует нескольких шагов, которые быть описаны в этой статье:
- Получение кадров изображения в режиме реального времени с веб-камеры. Я хочу, чтобы решение было найдено, наведя камеру на доску судоку. Это также накладывает некоторые требования к производительности, например, частота кадров видео не должна быть меньше 15 кадров в секунду.
- Нахождение контура доски судоку на каждом кадре. Есть готовые алгоритмы определения контуров, но на изображении может быть много лишних деталей, некоторые из них нежелательны и их нужно фильтровать.
- Извлечение однозначных изображений с доски и их распознавание. Нейронная сеть умеет распознавать цифры на изображении, но нам нужно подготовить для нее данные.
- Поиск решения для платы, если это возможно.
- Размещение цифр обратно на отображаемом изображении. То, что мы видим на экране, — это поток с веб-камеры. Чтобы сделать процесс более увлекательным, мы поместим недостающие цифры в те же места, где они должны быть на оригинальной доске.
Как найти доску судоку
Для первого теста я просто сделал скриншот случайной доски судоку со страницы Викимедиа:
Давайте загрузим изображение с помощью OpenCV:
import cv2 img = cv2.imread("sudoku.png")
Предварительная обработка состоит из двух шагов. Во-первых, мне нужно получить контур доски судоку. Для этого я использую методы Canny edgeDetection и метод findContours из OpenCV:
img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) edged = cv2.Canny(img_gray, 170, 490) contours, _ = cv2.findContours(edged, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
Результат будет примерно таким:
Контуры будут использоваться для поиска доски судоку. Но, как мы видим, есть много мелких контуров, которые мы должны игнорировать, мы добавим некоторые правила в код позже.
Во-вторых, я создаю монохромную копию изображения, которая будет использоваться для оптического распознавания символов:
blurred = cv2.GaussianBlur(img_gray, (11, 11), 0) img_bw = cv2.adaptiveThreshold(blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)
Для скриншотов это не критично, но важно для обработки изображений с камер прямых трансляций — с помощью метода adaptiveThreshold мы можем преобразовать изображения в монохромные, такие изображения более контрастны и лучше подходят для распознавания текста:
Теперь мы готовы обнаружить контуры на изображении. Доска судоку прямоугольная, поэтому мы проверяем только те контуры, которые не слишком маленькие, не слишком узкие и имеют 4 точки:
img_out = img.copy() w, h = img.shape[1], img.shape[0] for cntr in contours: imgx, imgy, imgw, imgh = cv2.boundingRect(cntr) if imgw < w/5 or imgw < h/5 or \ imgw/imgh < 0.25 or imgw/imgh > 1.5: continue # Approximate the contour with 4 points peri = cv2.arcLength(cntr, True) frm = cv2.approxPolyDP(cntr, 0.1*peri, True) if len(frm) != 4: continue # Converted image should fit into the original size board_size = max(imgw, imgh) if imgx + board_size >= w or imgy + board_size >= h: continue # Points should not be too close to each other # (use euclidian distance) if cv2.norm(frm[0][0] - frm[1][0], cv2.NORM_L2) < 0.1*peri or \ cv2.norm(frm[2][0] - frm[1][0], cv2.NORM_L2) < 0.1*peri or \ cv2.norm(frm[3][0] - frm[1][0], cv2.NORM_L2) < 0.1*peri or \ cv2.norm(frm[3][0] - frm[2][0], cv2.NORM_L2) < 0.1*peri: continue
Также оказалось, что некоторые контуры, возвращаемые OpenCV, имеют треугольную форму, но имеют 4 точки, 2 из которых расположены слишком близко друг к другу. Чтобы избежать этого, была добавлена дополнительная проверка евклидоварасстояния.
В целях отладки удобно рисовать контуры на выходном изображении и проверять результат:
# Draw sudoku contour using lines and points cv2.line(img_out, frm[0][0], frm[1][0], (0, 200, 0), thickness=3) cv2.line(img_out, frm[1][0], frm[2][0], (0, 200, 0), thickness=3) cv2.line(img_out, frm[2][0], frm[3][0], (0, 200, 0), thickness=3) cv2.line(img_out, frm[0][0], frm[3][0], (0, 200, 0), thickness=3) cv2.drawContours(img_out, frm, -1, (0, 255, 255), 10) cv2.imshow('image', img_out) cv2.waitKey(0)
Окончательный результат должен быть примерно таким:
При необходимости также можно сохранить изображение во временный файл:
cv2.imwrite("img_out.png", img_out)
Следующим важным шагом является применение преобразования перспективы. Это необязательно, но может значительно улучшить процесс распознавания искаженных или повернутых изображений, например:
В этом коде мы создаем преобразование перспективы из исходных координат контура в прямоугольную рамку:
def normalize_points(pts): rect = np.zeros((4, 2), dtype="float32") s = pts.sum(axis=1) rect[0] = pts[np.argmin(s)] rect[2] = pts[np.argmax(s)] diff = np.diff(pts, axis=1) rect[1] = pts[np.argmin(diff)] rect[3] = pts[np.argmax(diff)] return rect # Source and destination points for the perspective transform src_pts = normalize_points(frm.reshape((4, 2))) dst_pts = np.array([[0, 0], [board_size, 0], [board_size, board_size], [0, board_size]], dtype=np.float32) t_matrix = cv2.getPerspectiveTransform(approx__, dst_pts) _, t_matrix_inv = cv2.invert(t_matrix) # Convert images, colored and monochrome warped_disp = cv2.warpPerspective(img, t_matrix, (board_size, board_size)) warped_bw = cv2.warpPerspective(img_bw, t_matrix, (board_size, board_size))
Преобразование выполняется с помощью метода getPerspectiveTransform, для которого требуются два массива входных и выходных точек. Выяснилось, что иногда точки в контуре сохраняются слева направо, а иногда справа налево, для исправления этого был добавлен метод normalize_points — без этого метода выходное изображение может быть зеркальным. Используя это преобразование, я конвертирую два изображения — цветное для отображения и монохромное для процесса оптического распознавания.
Извлечение цифр
В конце предыдущего шага у нас есть доска для судоку квадратного размера, теперь мы можем разделить ее на 9 ячеек и извлечь все цифры. Подход в целом тот же — я получаю контуры каждой цифры и проверяю их размер, чтобы получить только цифры и пропустить нежелательную зернистость или шум:
images = [] cell_w, cell_h = board_size//9, board_size//9 for x in range(9): for y in range(9): x1, y1 = x*cell_w, y*cell_h x2, y2 = (x + 1)*cell_w, (y + 1)*cell_h cx, cy = (x1 + x2)//2, (y1 + y2)//2 w2, h2 = cell_w, cell_h # Find the contour of the digit crop = warped_bw[y1:y2, x1:x2] cntrs, _ = cv2.findContours(crop, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) for dc in cntrs: # w2, h2 = x2 - x1, y2 - y1 imgx2, imgy2, imgw2, imgh2 = cv2.boundingRect(dc) if 0.2*w2 < imgw2 < 0.8*w2 and 0.4*h2 < imgh2 < 0.8*h2: digit_img = crop[imgy2:imgy2 + imgh2, imgx2:imgx2 + imgw2] images.append((x, y, cx, cy, digit_img)) break
В результате должен получиться набор разделенных изображений, который выглядит так:
Теперь мы получили все изображения и готовы к следующему шагу — оптическому распознаванию.
OCR с помощью Tesseract
Мой первый подход заключался в использовании Tesseract, библиотеки с открытым исходным кодом для распознавания текста. Связать Tesseract с Python легко с помощью библиотеки pytesseract:
def predict_tesseract(images): results = [] for x, y, img_x, img_y, digit_img in images: value = predict_digit_tesseract(digit_img, x, y) results.append(value) return results def predict_digit_tesseract(digit_img, x, y): w, h = digit_img.shape if w > h: # Convert image to square size digit_img = cv2.copyMakeBorder(digit_img, 0, 0, (w - h)//2, w - h - (w - h)//2, cv2.BORDER_CONSTANT, value=(255,)) digit_img = cv2.copyMakeBorder(digit_img, w//10, w//10, w//10, w//10, cv2.BORDER_CONSTANT, value=(255,)) # Run OCR cf='-l eng --psm 8 --dpi 70 -c tessedit_char_whitelist=0123456789' res = pytesseract.image_to_string(digit_img, config=cf).strip() return int(res[0:1]) if len(res) > 0 else None
Может быть интересно разместить распознанные цифры на доске судоку. Для этого я использую инвертированную матрицу перспективного преобразования:
res = predict_tesseract(images) board = [0]*(9*9) for (x, y, img_x, img_y, digit_img), result in zip(images, res): if result: board[9*x + y] = result # Calculate coordinates on the original image orig = cv2.perspectiveTransform(np.array([[[img_x, img_y]]], dtype=np.float32), t_matrix_inv).reshape((2,)).astype(np.int32) cv2.putText(img_out, str(result), orig, cv2.FONT_HERSHEY_SIMPLEX, 1, (128, 0, 0), 2, cv2.LINE_AA, False)
После запуска кода получаем следующее:
Ну работает вроде. Результат неплохой, хотя расчет действительно медленный. Во-первых, Tesseract — довольно «тяжелая» библиотека, оптимизированная для распознавания отсканированных изображений. Во-вторых, в pyteserract нет пакетной обработки, поэтому невозможно отправить все цифры для OCR одним пакетом. В результате распознавание всех цифр с доски судоку заняло 2,1 секунды на процессоре Ryzen 9. И последнее, но не менее важное: Teserract не смог распознать цифру «9» в 3-й строке. Что ж, пришло время повеселиться и сделать собственное распознавание текста.
OCR с использованием PyTorch и CNN
Нейронная сеть
Я буду использовать сверточные нейронные сети (CNN), как мы знаем, CNN хорошо подходят для распознавания изображений. Те, кто не знаком с архитектурой CNN, могут прочитать хорошую статью об этом.
Давайте создадим модель нейронной сети:
import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim IMG_SIZE = 32 class Model(nn.Module): def __init__(self): super(Model, self).__init__() kernel_size = 5 self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=kernel_size, stride=1, padding=0) self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=kernel_size, stride=1, padding=0) out_layer_size = ((IMG_SIZE-kernel_size+1)-kernel_size+1)//2 self.dropout1 = nn.Dropout(0.25) self.dropout2 = nn.Dropout(0.5) self.fc1 = nn.Linear(64*out_layer_size*out_layer_size, 256) self.fc2 = nn.Linear(256, 10) def forward(self, x): x = self.conv1(x) x = F.relu(x) x = self.conv2(x) x = F.relu(x) x = F.max_pool2d(x, 2) x = self.dropout1(x) x = torch.flatten(x, 1) x = self.fc1(x) x = F.relu(x) x = self.dropout2(x) x = self.fc2(x) return F.log_softmax(x, dim=1)
Очевидно, что перед тем, как использовать модель, мы должны ее обучить, для этого нам нужен набор данных.
Набор данных
Распознавание цифр — простая задача для глубокого обучения, об этом есть много руководств. Но большинство авторов используют набор рукописных цифр MNIST:
Что неплохо для самообразования, но я никогда не видел доски судоку, напечатанной таким шрифтом. Итак, мы можем обучить нейронную сеть на этом наборе данных, но легко предсказать, что результаты будут не самыми лучшими.
К счастью, нам нужны только цифры от 0 до 9, расположенные в центре картинки. Легко создать класс набора данных, который будет генерировать разные цифры:
class DigitsDataset(torch.utils.data.Dataset): def __init__(self): # TTF files should be placed in the 'fonts' folder self.fonts = glob.glob("fonts/*.ttf") self.fonts_dict = {} self.digits_ = [None] * self.__len__() self.generate_all() def __len__(self): return 60000 def __getitem__(self, index): return self.digits_[index] def generate_all(self): print("Generating the digits dataset...") t_start = time.monotonic() for p in range(self.__len__()): if p % 10000 == 0: print(f" {p} of {self.__len__()}...") self.digits_[p] = self.generate_digit() print(f"Done, dT={time.monotonic() - t_start}s\n") def generate_digit(self): digit = random.randint(0, 9) data = self.generate_digit_pil(digit) return data, digit def generate_digit_pil(self, digit: int): text = str(digit) area_size = 2*IMG_SIZE img = Image.new("L", (area_size, area_size), (0,)) draw = ImageDraw.Draw(img) f_name, f_size = random.choice(self.fonts), random.randint(48, 64) key = f"{f_name}-{f_size}" if font_key not in self.fonts_dict: self.fonts_dict[key] = ImageFont.truetype(f_name, f_size) font = self.fonts_dict[key] text_x = area_size//2 + random.randint(-2, 2) text_y = area_size//2 - random.randint(-1, 1) draw.text((text_x, text_y), text, (255,), font=font, anchor="mm") transform = transforms.Compose([transforms.Resize([IMG_SIZE, IMG_SIZE]), transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]) resized = transform(img)[0].unsqueeze(0) return resized
Как видно из кода, я генерирую цифры, используя разные шрифты, разные размеры и позиции, что помогает лучше обучать сеть. Генерация цифр — медленный процесс, кэширование изображений увеличивает скорость обучения.
Мы можем легко проверить генерацию, отобразив цифры:
ds = DigitsDataset() images = [] for r in range(10): hor_images = [] for d in range(10): img = ds[10*r+d][0].reshape(IMG_SIZE, IMG_SIZE).detach().numpy() digit_img = cv2.copyMakeBorder(img, 2, 2, 2, 2, cv2.BORDER_CONSTANT, value=(128,)) hor_images.append(digit_img) images.append(np.concatenate(hor_images, axis=1)) cv2.imshow("Dataset", np.concatenate(images, axis=0)) cv2.waitKey(0)
Видно, что результаты выглядят лучше, чем рукописные цифры:
Обучение модели
Теперь у нас есть набор данных, и мы можем обучить модель:
import time max_epochs = 12 batch_size_train = 10 log_interval = 600 model_file_name = "ocr_model.pt" use_cuda = torch.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") # Dataset instance dataset = DigitsDataset() train_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size_train, shuffle=True) # Create the model model = Model().to(device) # Train model.train() print("Training the model, use CUDA:", use_cuda) optimizer = optim.Adadelta(model.parameters(), lr=0.01) for epoch in range(max_epochs): print(f"Epoch: {epoch + 1} of {max_epochs}") t1 = time.monotonic() total_loss = 0 for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() output = model(data) loss = F.nll_loss(output, target) loss.backward() optimizer.step() if batch_idx % log_interval == 0: print(' Train [{}/{} ({:.0f}%)] Loss: {:.6f}'.format(batch_idx * len(data), len(train_loader.dataset), 100 * batch_idx / len(train_loader), loss.item())) total_loss += loss.item() print(f" Total loss: {total_loss}, dT={time.monotonic() - t1}s") torch.save(model.state_dict(), model_file_name) print("Model saved to", model_file_name)
Каждая эпоха обучения требует около 12 секунд при использовании CUDA и 80 секунд при использовании ЦП, поэтому настоятельно рекомендуется хорошая видеокарта. Если нет доступного графического процессора с поддержкой CUDA (здесь я говорю привет пользователям Mac;), Google Colab — хорошая бесплатная альтернатива, он может поддерживать графический процессор и работать довольно быстро. Чтобы использовать внешние файлы в блокноте Colab, мы можем поместить файлы (исходный код Python и шрифты TTF) на Google Диск и смонтировать этот диск, используя код:
from google.colab import drive drive.mount('/content/drive')
После этого папка Google Диска будет доступна из кода (пути к файлам также должны быть изменены на что-то вроде /content/drive/MyDrive/Colab Notebooks/…). И это работает — после включения GPU в настройках Google Colab эпоха обучения заняла 21 секунду. Он в 4 раза быстрее, чем мой 16-ядерный процессор, и всего в 2 раза медленнее, чем графический процессор, что на удивление неплохо, учитывая, что сервис Google Colab бесплатный, а я заплатил за свой графический процессор около 1200 долларов.
Признание
На предыдущем шаге был создан «ocr_model.pt», теперь мы можем использовать обученную модель для распознавания цифр. Во-первых, давайте загрузим модель:
model_file_name = "ocr_model.pt" model = Model() device = "cpu" model.load_state_dict(torch.load(model_file_name, map_location=torch.device(device)))
Как мы видим, устройство настроено на «ЦП» — процесс распознавания быстрый, и мне здесь не нужен графический процессор. На самом деле инициализация всех библиотек графического процессора и обмен данными могут занять больше времени, чем распознавание на базе процессора.
У меня уже есть метод predict_tesseract, я создам аналогичный метод для прогнозирования цифр, используя нашу модель PyTorch:
def predict_pytorch(model: Model, images: List): transform = transforms.Compose([transforms.ToPILImage(), transforms.Resize([IMG_SIZE, IMG_SIZE]), transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]) # Prepare images for the recognition images_ = [] for x, y, img_x, img_y, digit_img in images: w, h = digit_img.shape # Convert image to square size if w > h: img_square = cv2.copyMakeBorder(digit_img, 10, 10, 10 + (w - h)//2, 10 + w - h - (w - h)//2, cv2.BORDER_CONSTANT, value=(255,)) else: img_square = cv2.copyMakeBorder(digit_img, 10 + (h - w)//2, 10 + h - w - (h - w)//2, 10, 10, cv2.BORDER_CONSTANT, value=(255,)) data = transform(~img_square).unsqueeze(0) images_.append(data) if len(images_) == 0: return [] # Convert separated images to the single Pytorch tensor data = torch.cat(images_) # Run OCR model model.eval() with torch.no_grad(): out = model(data.to(device)) p = out.data.max(1, keepdim=True)[1].reshape((len(images_), )) return p.tolist()
Теперь нам нужно только заменить старый метод на новый, больше никаких изменений в коде не требуется. После этого мы можем получить результаты:
Это очень хорошо — все цифры распознались, а общее время — всего 0,01 с вместо 2,1 с для Tesseract!
Кстати, для тех, кто хочет протестировать набор данных MNIST на предмет распознавания, сделать это несложно. Требуется только заменить экземпляр DigitsDataset() в обучающем коде на класс datasets.MNIST:
dataset = datasets.MNIST(mnist_folder, train=True, download=True, transform=transforms.Compose( [transforms.Resize([IMG_SIZE, IMG_SIZE]), transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]))
Мы можем визуализировать и использовать набор данных таким же образом, никаких изменений кода не потребуется. И интересно отметить, что результаты, на удивление, не такие уж плохие, большинство цифр все еще можно распознать правильно. Вот результат той же нейронной сети, обученной MNIST:
Как мы видим, только одна из цифр была распознана неправильно, вместо «6» была распознана цифра «9».
Решение судоку
Мы почти готовы выполнить задание — мы узнали все цифры и теперь можем найти решение для доски судоку. Я не эксперт в судоку, поэтому идею алгоритма я нашел в проекте Sudoku-GUI-Solver. Это работает, вроде. Код использует алгоритм возврата для поиска решения, но проблема в том, что Python хорош для высокоуровневых операций с данными, но не так хорош для обработки чисел — время обработки составляло около 10 секунд на доску, что абсолютно далеко от реального времени! Простое решение — использовать Numba — компилятор Python. Технически это легко исправить, в большинстве случаев достаточно добавить директиву @njit в определение метода. И это работает, время обработки уменьшилось с 10 до 0,8с — что само по себе хорошо, но все же недостаточно для нашей задачи.
Более эффективный способ — это переписать алгоритм на C, что я собственно и сделал, сам код имеет всего около 50 строк кода. C-функции можно вызывать из Python с помощью ctypes:
import ctypes lib = None def solve_c(bo: List) -> bool: global lib if lib is None: lib = ctypes.CDLL('solver_lib.so') lib.solve.argtypes = [ctypes.POINTER(ctypes.c_int)] board_data = (ctypes.c_int * len(bo))(*bo) res = lib.solve(board_data) if res: bo[:] = list(board_data) return res board = [1, 6, 0, 0, .... ] res = solve_c(board) print("Solution found:", res)
Как мы видим, я использую ctypes для преобразования списка Python в C-массив и извлечения данных обратно после вычисления. C-файл содержит всю «магию»:
DLL_EXPORT int solve(int *bo) { ... }
Единственный минус использования C-binding в том, что перед использованием программы необходимо скомпилировать библиотеку под нужную платформу:
# Linux: gcc -shared -Wl,-soname,solver_lib -o solver_lib.so -fPIC solver_lib.c # OSX: gcc -shared -Wl,-install_name,solver_lib.so -o solver_lib.so -fPIC solver_lib.c # Windows: should be doable via the Visual Studio DLL Project
И теперь со скоростью все в порядке — время расчета составляет около 0,01 с, и мы готовы к работе в реальном времени.
В ходе тестирования было добавлено еще одно изменение в алгоритм. Оказалось, что исходный алгоритм ожидает на входе правильные значения, что не всегда верно для OCR. Иногда цифра 7 может быть неправильно распознана как 1 или 8 как 0, и алгоритм может просто застрять в бесконечном цикле, если он имеет неправильные входные значения. Добавлена простая проверка, которая проверяет, что во входных данных нет повторяющихся чисел, если такие числа будут найдены, то доску, очевидно, решить не удастся.
Переход в реальном времени
Наш последний шаг — объединить все части вместе и запустить этот код с потоком камеры. Это легко сделать с помощью OpenCV:
model = Model() device = "cpu" model_file_name = "ocr_model.pt" model.load_state_dict(torch.load(model_file_name, map_location=torch.device(device))) model.eval() cap = cv2.VideoCapture(0) width, height, fps = 1280, 720, 15 cap.set(cv2.CAP_PROP_FRAME_WIDTH, width) cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height) cap.set(cv2.CAP_PROP_FPS, fps) print(f"Starting the video stream: {width}x{height}") while True: # Capture images frame-by-frame ret, frame = cap.read() frame_orig = frame.copy() try: # Process res, img_out = process_image(model, frame) # Display cv2.imshow('Frame', img_out) except Exception as e: # If something was wrong, save frame for debugging cv2.imwrite("crash.png", frame_orig) break # Process key codes key_code = cv2.waitKey(1) if key_code & 0xFF == ESC_KEY: break cap.release() cv2.destroyAllWindows()
Как мы видим, кода относительно немного, мы получаем кадры с веб-камеры и отправляем их в метод process_image, содержащий всю описанную ранее логику. Я также добавил предварительный просмотр доски судоку, который помогает оценить качество преобразования перспективы и извлечения цифр. Анимация показывает, что конечный результат довольно точен:
Заключение
Работа с данными и нейронными сетями может быть увлекательной. И как мы видим, само по себе распознавание цифр — относительно простая и известная задача, но объединение всего воедино и запуск «в продакшн» может потребовать массу дополнительных шагов, оптимизаций и доработок. В любом случае, моя первоначальная оценка, что такое задание можно выполнить за 2–4 часа, была ужасно ошибочной, но это делает задачу еще более увлекательной.
Желающие проводить эксперименты самостоятельно могут скачать код с GitHub.
Спасибо за прочтение.