sudoku_puzzle_extractor.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9f61cadb-e562-4bcd-8621-6868797f0bcc",
   "metadata": {},
   "source": [
    "## Sudoku AI Solver\n",
    "\n",
    "In this notebook, a sudoku puzzle is extracted from an image using OpenCV and Deep Learning, then solved using a Python script. \n",
    "\n",
    "Refer to the main documentation at: https://thenoobinventor.github.io/sudoku-ai-solver/ for more information on the different steps taken in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b32b6eb0-3a4f-4a33-aad5-2d1d8129ee12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Disable tensorflow warning messages\n",
    "\n",
    "import os\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "307796e3-3159-449d-a7ba-5f530977de93",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import imutils\n",
    "import copy\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from keras.models import load_model\n",
    "from tensorflow.keras.preprocessing.image import img_to_array\n",
    "from solve_sudoku import *\n",
    "from imutils.perspective import four_point_transform\n",
    "from skimage.segmentation import clear_border"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfe308b4-02d2-4753-bdeb-3e77ed23fdb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_puzzle(img, debug=False):\n",
    "    \"\"\"\n",
    "    Extract sudoku puzzle from an image.\n",
    "    \"\"\"\n",
    "    \n",
    "    # Convert image to grayscale and blur it to filter out noise\n",
    "    gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)\n",
    "    blurred = cv2.GaussianBlur(gray, (7,7), 3)\n",
    "\n",
    "    # Apply adaptive thresholding\n",
    "    thresh = cv2.adaptiveThreshold(blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 2)\n",
    "\n",
    "    # Find contours in the thresholded image and sort them by size in descending order\n",
    "    cnts = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n",
    "    cnts = imutils.grab_contours(cnts)\n",
    "    cnts = sorted(cnts, key=cv2.contourArea, reverse=True)\n",
    "    \n",
    "    # Initialize a contour that corresponds to the puzzle outline\n",
    "    puzzle_cnt = None\n",
    "    \n",
    "    # Loop over the contours\n",
    "    for c in cnts:\n",
    "        # Approximate the coutour\n",
    "        peri = cv2.arcLength(c, True)\n",
    "        approx = cv2.approxPolyDP(c, 0.02 * peri, True)\n",
    "\n",
    "        # If the approximated contour has 4 points, then it is assumed that \n",
    "        # this contour is the puzzle outline\n",
    "        if len(approx) == 4:\n",
    "            puzzle_cnt = approx\n",
    "            break\n",
    "        \n",
    "    puzzle_outline = img.copy()\n",
    "    cv2.drawContours(puzzle_outline, [puzzle_cnt], -1, (0, 255, 0), 3)\n",
    "\n",
    "    # Apply a four point perspective transform to both the original image and \n",
    "    # grayscale image to obtain a top-down bird's eye view of the puzzle\n",
    "    color_puzzle = four_point_transform(img, puzzle_cnt.reshape(4,2)) \n",
    "    gray_puzzle = four_point_transform(gray, puzzle_cnt.reshape(4,2))  \n",
    "    \n",
    "    if debug:\n",
    "        display_img(img, 'Original Image')\n",
    "        display_img(thresh, 'Thresholded Image')\n",
    "        display_img(puzzle_outline, 'Puzzle Outline')\n",
    "        display_img(gray_puzzle, 'Extracted Grayscale Puzzle')\n",
    "        \n",
    "    return color_puzzle, gray_puzzle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a681f4c-1c29-44b3-8a73-b90ac797d100",
   "metadata": {},
   "outputs": [],
   "source": [
    "def display_img(img, title):\n",
    "    \"\"\"\n",
    "    Function to display an image with a specified title.\n",
    "    \"\"\"\n",
    "    \n",
    "    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "    \n",
    "    fig = plt.figure(figsize=(12,10))\n",
    "    ax = fig.add_subplot(111)\n",
    "    ax.set_title(title, fontsize=22)\n",
    "    ax.imshow(img,cmap='gray')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6667a7e-4592-4f2a-8127-fc4a8ca9f1ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_cell_locations(step_x, step_y):\n",
    "    \"\"\"\n",
    "    Generate (x, y) coordinates for each cell of the sudoku board.\n",
    "    \"\"\"\n",
    "\n",
    "    # Initialize a list to store (x,y) coordinates of each cell location\n",
    "    cell_locs = []\n",
    "\n",
    "    # Loop over the grid lcoations\n",
    "    for y in range(9):\n",
    "        # Initialize the current list of cell locations\n",
    "        row = []\n",
    "\n",
    "        for x in range(9):\n",
    "            # Compute the starting and ending (x,y) coordinates of the current cell\n",
    "            start_x = x * step_x\n",
    "            start_y = y * step_y\n",
    "            end_x = (x + 1) * step_x\n",
    "            end_y = (y + 1) * step_y\n",
    "\n",
    "            # Add the (x,y) coordinates to the cell locations list\n",
    "            row.append((start_x, start_y, end_x, end_y))\n",
    "\n",
    "            # Crop the cell from the gray_puzzle transformed image and then extract\n",
    "            # the digit from the cell\n",
    "            cell = gray_puzzle[start_y:end_y, start_x:end_x]\n",
    "            digit = extract_digit(cell)\n",
    "\n",
    "            # Confirm that the digit is not empty\n",
    "            if digit is not None:\n",
    "                classify_digit(digit, x, y)\n",
    "\n",
    "        # Add the row to the cell locations\n",
    "        cell_locs.append(row)\n",
    "\n",
    "    return cell_locs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39d12ea5-7adc-4716-8db1-e8843c304df3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_digit(cell):\n",
    "    \"\"\"\n",
    "    Extracts and returns the digit of the passed in cell. If no digit was identified, the function \n",
    "    returns None.\n",
    "    \"\"\"\n",
    "    \n",
    "    # Apply automatic thresholding to the cell and then clear any connected borders \n",
    "    # that touch the border of the cell\n",
    "    thresh = cv2.threshold(cell, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)[1]\n",
    "    thresh = clear_border(thresh)\n",
    "    \n",
    "    # Find contours in the thresholded cell\n",
    "    cnts = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n",
    "    cnts = imutils.grab_contours(cnts)\n",
    "    \n",
    "    # If no contours were found then this is an empty cell\n",
    "    if len(cnts)==0:\n",
    "        return None\n",
    "    \n",
    "    # Otherwise find the largest contour in the cell and create a mask for the contour\n",
    "    c = max(cnts, key=cv2.contourArea)\n",
    "    mask = np.zeros(thresh.shape, dtype=\"uint8\")\n",
    "    cv2.drawContours(mask, [c], -1, 255, -1)\n",
    "    \n",
    "    # Compute the perecentage of masked pixels relative to the total area of the image\n",
    "    (h, w) = thresh.shape\n",
    "    percent_filled = cv2.countNonZero(mask) / float(w * h)\n",
    "    \n",
    "    # If less than 3% of the mask is filled then we are looking at noise and can safely ignore the contour\n",
    "    if percent_filled < 0.03:\n",
    "        return None\n",
    "    \n",
    "    # Apply the mask to the thresholded cell\n",
    "    digit = cv2.bitwise_and(thresh, thresh, mask=mask)\n",
    "\n",
    "    return digit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51583768",
   "metadata": {},
   "outputs": [],
   "source": [
    "def classify_digit(digit, x, y):\n",
    "    \"\"\"\n",
    "    Classifies the passed in digit using the OCR model then assigns the digit to \n",
    "    corresponding cell of the unsolved sudoku board.\n",
    "    \"\"\"\n",
    "    \n",
    "    global unsolved_board\n",
    "    \n",
    "    # Resize the cell to 28x28 pixels and prepare it classification \n",
    "    # 28x28 is the size of images in the MNIST dataset\n",
    "    roi = cv2.resize(digit, (28, 28))\n",
    "    roi = roi.astype(\"float\")/255.0\n",
    "    roi = img_to_array(roi)\n",
    "    roi = np.expand_dims(roi, axis=0)\n",
    "\n",
    "    # Classify the digit and update the sudoku board with the prediction\n",
    "    pred = model.predict(roi, verbose=0).argmax(axis=1)[0]\n",
    "    unsolved_board[y, x] = pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7733677-e025-4986-8901-b3de612630a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def display_solutions(cell_locs, color_puzzle):\n",
    "    \"\"\"\n",
    "    This function displays the solutions of the sudoku puzzle on the color puzzle image.\n",
    "    \"\"\"\n",
    "    \n",
    "    # Loop over the cell locations and boards\n",
    "    for (cell_row, unsolved_board_row, solved_board_row) in zip(cell_locs, unsolved_board, solved_board):\n",
    "        \n",
    "        # Loop over individual cells in the row\n",
    "        for (box, unsolved_digit, solved_digit) in zip(cell_row, unsolved_board_row, solved_board_row):\n",
    "            if unsolved_digit == 0:\n",
    "                # Unpack the cell coordinates\n",
    "                start_x, start_y, end_x, end_y = box\n",
    "\n",
    "                # Compute the coordinates of where the digit will be drawn on the output puzzle image\n",
    "                text_x = int((end_x - start_x) * 0.33)\n",
    "                text_y = int((end_y - start_y) * -0.2)\n",
    "                text_x += start_x\n",
    "                text_y += end_y\n",
    "\n",
    "                # Draw the digit on the sudoku puzzle image\n",
    "                cv2.putText(\n",
    "                    color_puzzle, str(solved_digit), (text_x, text_y), \n",
    "                    cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 255), 2)\n",
    "    \n",
    "    display_img(color_puzzle, 'Solved Puzzle')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ba8c4c7-544b-46ce-ac39-6851a004f749",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Main Program Cell\n",
    "\n",
    "# Load image; second image is sudoku2.jpg\n",
    "img = cv2.imread('sudoku_images/sudoku.jpg')\n",
    "img = imutils.resize(img, width=600)\n",
    "\n",
    "# Find puzzle in the image; set debug to False to disable displaying image processing steps\n",
    "color_puzzle, gray_puzzle = find_puzzle(img, debug=True)\n",
    "\n",
    "# Initialize sudoku board\n",
    "unsolved_board = np.zeros((9,9), dtype='int')\n",
    "\n",
    "# Sudoku is a 9x9 grid (81 individual cells), location of each cell can be inferred by\n",
    "# dividing the gray_puzzle image into a 9x9 grid\n",
    "step_x = gray_puzzle.shape[1] // 9\n",
    "step_y = gray_puzzle.shape[0] // 9\n",
    "\n",
    "# Load model to detect digits\n",
    "model = load_model('model_files/digit_classifier_model.h5')\n",
    "\n",
    "# List of the (x-y) coordinate location of each cell\n",
    "cell_locs = generate_cell_locations(step_x, step_y)\n",
    "\n",
    "# Convert board to a list before solving the puzzle\n",
    "unsolved_board = unsolved_board.tolist()\n",
    "\n",
    "# Create a deep copy of the unsolved_board to be passed into the solve_puzzle() function\n",
    "unsolved_board_copy = copy.deepcopy(unsolved_board)\n",
    "\n",
    "# Note: debug should only be set to True if there are any errors encountered in solving the puzzle. The debug flag helps to identify \n",
    "# where the error might be such as puzzle elements not being integers or presence of duplicates in a row, column, 3x3 block. \n",
    "# However, it results in a significantly slower performance of the function.\n",
    "#\n",
    "# The error checks enabled by the debug flag were also used in unit testing the solve_puzzle() function.\n",
    "\n",
    "# Solve puzzle\n",
    "solved_board = solve_puzzle(unsolved_board_copy, debug=False)\n",
    "\n",
    "# Display puzzle solutions\n",
    "if not solved_board:\n",
    "    print('\\nPuzzle could not be solved, check training model.\\n')\n",
    "else:\n",
    "    display_solutions(cell_locs, color_puzzle)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}