Однажды я решил проверить, можно ли решать судоку в реальном времени с помощью машинного обучения и обработки изображений. Я изначально думал, что это можно сделать за 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.

Спасибо за прочтение.