chess_board_recognizer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'\n",
    "\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "# Imports for visualization\n",
    "import PIL.Image\n",
    "from io import BytesIO\n",
    "from IPython.display import clear_output, Image, display\n",
    "import scipy.ndimage as nd\n",
    "import scipy.signal\n",
    "\n",
    "\n",
    "def display_array(a, format='jpeg', rng=[0, 1]):\n",
    "    a = (a - rng[0]) / float(rng[1] - rng[0]) * 255\n",
    "    a = np.uint8(np.clip(a, 0, 255))\n",
    "    file = BytesIO()\n",
    "    PIL.Image.fromarray(a).save(file, format)\n",
    "    display(Image(data=file.getvalue()))\n",
    "\n",
    "def ImageName(str):\n",
    "    img_file = str\n",
    "\n",
    "def make_kernel(a):\n",
    "    \"\"\"Transform a 2D array into a convolution kernel\"\"\"\n",
    "    a = np.asarray(a)\n",
    "    a = a.reshape(list(a.shape) + [1, 1])\n",
    "    return tf.constant(a, dtype=1)\n",
    "\n",
    "\n",
    "def simple_conv(x, k):\n",
    "    \"\"\"A simplified 2D convolution operation\"\"\"\n",
    "    x = tf.expand_dims(tf.expand_dims(x, 0), -1)\n",
    "    y = tf.nn.depthwise_conv2d(x, k, [1, 1, 1, 1], padding='SAME')\n",
    "    return y[0, :, :, 0]\n",
    "\n",
    "\n",
    "def gradientx(x):\n",
    "    \"\"\"Compute the x gradient of an array\"\"\"\n",
    "    gradient_x = make_kernel([[-1., 0., 1.],\n",
    "                              [-1., 0., 1.],\n",
    "                              [-1., 0., 1.]])\n",
    "    return simple_conv(x, gradient_x)\n",
    "\n",
    "\n",
    "def gradienty(x):\n",
    "    \"\"\"Compute the x gradient of an array\"\"\"\n",
    "    gradient_y = make_kernel([[-1., -1, -1], [0., 0, 0], [1., 1, 1]])\n",
    "    return simple_conv(x, gradient_y)\n",
    "\n",
    "\n",
    "def corners(x):\n",
    "    \"\"\"Find chess square corners in an array\"\"\"\n",
    "    chess_corner = make_kernel([[-1., 0, 1], [0., 0., 0.], [1., 0, -1]])\n",
    "    return simple_conv(x, chess_corner)\n",
    "\n",
    "\n",
    "def dilate(x, size=3):\n",
    "    \"\"\"Dilate\"\"\"\n",
    "    kernel = make_kernel(np.ones([size, size], dtype=np.float32))\n",
    "    return tf.clip_by_value(simple_conv(x, kernel),\n",
    "                            np.float32(1),\n",
    "                            np.float32(2)) - np.float32(1)\n",
    "\n",
    "\n",
    "def erode(x, size=3):\n",
    "    \"\"\"Erode\"\"\"\n",
    "    kernel = make_kernel(np.ones([size, size]))\n",
    "    return tf.clip_by_value(simple_conv(x, kernel),\n",
    "                            np.float32(size * size - 1),\n",
    "                            np.float32(size * size)) - np.float32(size * size - 1)\n",
    "\n",
    "\n",
    "def opening(x, size=3):\n",
    "    return dilate(erode(x, size), size)\n",
    "\n",
    "\n",
    "def closing(x, size=3):\n",
    "    return erode(dilate(x, size), size)\n",
    "\n",
    "\n",
    "def skeleton(x, size=3):\n",
    "    \"\"\"Skeletonize\"\"\"\n",
    "    return tf.clip_by_value(erode(x) - opening(erode(x)),\n",
    "                            0.,\n",
    "                            1.)\n",
    "\n",
    "\n",
    "# Get our grayscale image matrix\n",
    "\n",
    "\n",
    "def checkMatch(lineset):\n",
    "    \"\"\"Checks whether there exists 7 lines of consistent increasing order in set of lines\"\"\"\n",
    "    linediff = np.diff(lineset)\n",
    "    x = 0\n",
    "    cnt = 0\n",
    "    for line in linediff:\n",
    "        # Within 5 px of the other (allowing for minor image errors)\n",
    "        if np.abs(line - x) < 5:\n",
    "            cnt += 1\n",
    "        else:\n",
    "            cnt = 0\n",
    "            x = line\n",
    "    return cnt == 5\n",
    "\n",
    "\n",
    "def pruneLines(lineset):\n",
    "    \"\"\"Prunes a set of lines to 7 in consistent increasing order (chessboard)\"\"\"\n",
    "    linediff = np.diff(lineset)\n",
    "    x = 0\n",
    "    cnt = 0\n",
    "    start_pos = 0\n",
    "    for i, line in enumerate(linediff):\n",
    "        # Within 5 px of the other (allowing for minor image errors)\n",
    "        if np.abs(line - x) < 5:\n",
    "            cnt += 1\n",
    "            if cnt == 5:\n",
    "                end_pos = i + 2\n",
    "                return lineset[start_pos:end_pos]\n",
    "        else:\n",
    "            cnt = 0\n",
    "            x = line\n",
    "            print (i, x)\n",
    "            start_pos = i\n",
    "    return lineset\n",
    "\n",
    "\n",
    "def skeletonize_1d(arr):\n",
    "    \"\"\"return skeletonized 1d array (thin to single value, favor to the right)\"\"\"\n",
    "    _arr = arr.copy()  # create a copy of array to modify without destroying original\n",
    "    # Go forwards\n",
    "    for i in range(_arr.size - 1):\n",
    "        # Will right-shift if they are the same\n",
    "        if arr[i] <= _arr[i + 1]:\n",
    "            _arr[i] = 0\n",
    "\n",
    "    # Go reverse\n",
    "    for i in np.arange(_arr.size - 1, 0, -1):\n",
    "        if _arr[i - 1] > _arr[i]:\n",
    "            _arr[i] = 0\n",
    "    return _arr\n",
    "\n",
    "\n",
    "def getChessLines(hdx, hdy, hdx_thresh, hdy_thresh):\n",
    "    \"\"\"Returns pixel indices for the 7 internal chess lines in x and y axes\"\"\"\n",
    "    # Blur\n",
    "    gausswin = scipy.signal.gaussian(21, 4)\n",
    "    gausswin /= np.sum(gausswin)\n",
    "\n",
    "    # Blur where there is a strong horizontal or vertical line (binarize)\n",
    "    blur_x = np.convolve(hdx > hdx_thresh, gausswin, mode='same')\n",
    "    blur_y = np.convolve(hdy > hdy_thresh, gausswin, mode='same')\n",
    "\n",
    "    skel_x = skeletonize_1d(blur_x)\n",
    "    skel_y = skeletonize_1d(blur_y)\n",
    "\n",
    "    # Find points on skeletonized arrays (where returns 1-length tuple)\n",
    "    lines_x = np.where(skel_x)[0]  # vertical lines\n",
    "    lines_y = np.where(skel_y)[0]  # horizontal lines\n",
    "\n",
    "    # Prune inconsistent lines\n",
    "    lines_x = pruneLines(lines_x)\n",
    "    lines_y = pruneLines(lines_y)\n",
    "\n",
    "    is_match = len(lines_x) == 7 and len(\n",
    "        lines_y) == 7 and checkMatch(lines_x) and checkMatch(lines_y)\n",
    "\n",
    "    return lines_x, lines_y, is_match\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "def getChessTiles(a, lines_x, lines_y):\n",
    "    \"\"\"Split up input grayscale array into 64 tiles stacked in a 3D matrix using the chess linesets\"\"\"\n",
    "    # Find average square size, round to a whole pixel for determining edge pieces sizes\n",
    "    stepx = np.int32(np.round(np.mean(np.diff(lines_x))))\n",
    "    stepy = np.int32(np.round(np.mean(np.diff(lines_y))))\n",
    "\n",
    "    padr_x = 0\n",
    "    padl_x = 0\n",
    "    padr_y = 0\n",
    "    padl_y = 0\n",
    "\n",
    "    if lines_x[0] - stepx < 0:\n",
    "        padl_x = np.abs(lines_x[0] - stepx)\n",
    "    if lines_x[-1] + stepx > a.shape[1] - 1:\n",
    "        padr_x = np.abs(lines_x[-1] + stepx - a.shape[1])\n",
    "    if lines_y[0] - stepy < 0:\n",
    "        padl_y = np.abs(lines_y[0] - stepy)\n",
    "    if lines_y[-1] + stepx > a.shape[0] - 1:\n",
    "        padr_y = np.abs(lines_y[-1] + stepy - a.shape[0])\n",
    "\n",
    "    # New padded array\n",
    "#     print \"Padded image to\", ((padl_y,padr_y),(padl_x,padr_x))\n",
    "    a2 = np.pad(a, ((padl_y, padr_y), (padl_x, padr_x)), mode='edge')\n",
    "\n",
    "    setsx = np.hstack([lines_x[0] - stepx, lines_x,\n",
    "                       lines_x[-1] + stepx]) + padl_x\n",
    "    setsy = np.hstack([lines_y[0] - stepy, lines_y,\n",
    "                       lines_y[-1] + stepy]) + padl_y\n",
    "\n",
    "    a2 = a2[setsy[0]:setsy[-1], setsx[0]:setsx[-1]]\n",
    "    setsx -= setsx[0]\n",
    "    setsy -= setsy[0]\n",
    "#     display_array(a2, rng=[0,255])\n",
    "#     print \"X:\",setsx\n",
    "#     print \"Y:\",setsy\n",
    "\n",
    "    # Matrix to hold images of individual squares (in grayscale)\n",
    "#     print \"Square size: [%g, %g]\" % (stepy, stepx)\n",
    "    squares = np.zeros([np.round(stepy), np.round(stepx), 64], dtype=np.uint8)\n",
    "\n",
    "    # For each row\n",
    "    for i in range(0, 8):\n",
    "        # For each column\n",
    "        for j in range(0, 8):\n",
    "            # Vertical lines\n",
    "            x1 = setsx[i]\n",
    "            x2 = setsx[i + 1]\n",
    "            padr_x = 0\n",
    "            padl_x = 0\n",
    "            padr_y = 0\n",
    "            padl_y = 0\n",
    "\n",
    "            if (x2 - x1) > stepx:\n",
    "                if i == 7:\n",
    "                    x1 = x2 - stepx\n",
    "                else:\n",
    "                    x2 = x1 + stepx\n",
    "            elif (x2 - x1) < stepx:\n",
    "                if i == 7:\n",
    "                    # right side, pad right\n",
    "                    padr_x = stepx - (x2 - x1)\n",
    "                else:\n",
    "                    # left side, pad left\n",
    "                    padl_x = stepx - (x2 - x1)\n",
    "            # Horizontal lines\n",
    "            y1 = setsy[j]\n",
    "            y2 = setsy[j + 1]\n",
    "\n",
    "            if (y2 - y1) > stepy:\n",
    "                if j == 7:\n",
    "                    y1 = y2 - stepy\n",
    "                else:\n",
    "                    y2 = y1 + stepy\n",
    "            elif (y2 - y1) < stepy:\n",
    "                if j == 7:\n",
    "                    # right side, pad right\n",
    "                    padr_y = stepy - (y2 - y1)\n",
    "                else:\n",
    "                    # left side, pad left\n",
    "                    padl_y = stepy - (y2 - y1)\n",
    "            # slicing a, rows sliced with horizontal lines, cols by vertical lines so reversed\n",
    "            # Also, change order so its A1,B1...H8 for a white-aligned board\n",
    "            # Apply padding as defined previously to fit minor pixel offsets\n",
    "            squares[:, :, (7 - j) * 8 + i] = np.pad(a2[y1:y2, x1:x2],\n",
    "                                                    ((padl_y, padr_y), (padl_x, padr_x)), mode='edge')\n",
    "    return squares\n",
    "\n",
    "\n",
    "\n",
    "def generateTileset(str, folder_name, tiles_directory):\n",
    "    np.set_printoptions(suppress=True)\n",
    "\n",
    "    sess = tf.InteractiveSession()\n",
    "    img_file = str\n",
    "    folder = folder_name\n",
    "\n",
    "    img = PIL.Image.open(\"%s/%s\" % (folder, img_file))\n",
    "\n",
    "    #print (\"Loaded %s (%dpx x %dpx)\" % \\\n",
    "    #    (img_file, img.size[0], img.size[1]))\n",
    "\n",
    "    # Resize if image larger than 2k pixels on a side\n",
    "    if img.size[0] > 2000 or img.size[1] > 2000:\n",
    "        print (\"Image too big (%d x %d)\" % (img.size[0], img.size[1]))\n",
    "        new_size = 500.0  # px\n",
    "        if img.size[0] > img.size[1]:\n",
    "            # resize by width to new limit\n",
    "            ratio = new_size / img.size[0]\n",
    "        else:\n",
    "            # resize by height\n",
    "            ratio = new_size / img.size[1]\n",
    "        print (\"Reducing by factor of %.2g\" % (1. / ratio))\n",
    "        img = img.resize(img.size * ratio, PIL.Image.ADAPTIVE)\n",
    "        print (\"New size: (%d x %d)\" % (img.size[0], img.size[1]))\n",
    "\n",
    "    # See original image\n",
    "    #display_array(np.asarray(img), rng=[0, 255])\n",
    "\n",
    "    # Convert to grayscale and array\n",
    "    a = np.asarray(img.convert(\"L\"), dtype=np.float32)\n",
    "\n",
    "    # Display array\n",
    "    #display_array(a, rng=[0, 255])\n",
    "    A = tf.Variable(a)\n",
    "\n",
    "\n",
    "    Dx = gradientx(A)\n",
    "    Dy = gradienty(A)\n",
    "\n",
    "\n",
    "    tf.global_variables_initializer().run()\n",
    "\n",
    "    #display_array(Dx.eval(), rng=[-255, 255])\n",
    "    #display_array(Dy.eval(), rng=[-255, 255])\n",
    "\n",
    "    Dx_pos = tf.clip_by_value(Dx, 0., 255., name=\"dx_positive\")\n",
    "    Dx_neg = tf.clip_by_value(Dx, -255., 0., name='dx_negative')\n",
    "    Dy_pos = tf.clip_by_value(Dy, 0., 255., name=\"dy_positive\")\n",
    "    Dy_neg = tf.clip_by_value(Dy, -255., 0., name='dy_negative')\n",
    "\n",
    "\n",
    "    hough_Dx = tf.reduce_sum(Dx_pos, 0) * \\\n",
    "        tf.reduce_sum(-Dx_neg, 0) / (a.shape[0] * a.shape[0])\n",
    "    hough_Dy = tf.reduce_sum(Dy_pos, 1) * \\\n",
    "        tf.reduce_sum(-Dy_neg, 1) / (a.shape[1] * a.shape[1])\n",
    "\n",
    "\n",
    "    import matplotlib.pyplot as plt\n",
    "\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2, sharey=True, figsize=(15, 5))\n",
    "\n",
    "    # Arbitrarily choose half of max value as threshold, since they're such strong responses\n",
    "    hough_Dx_thresh = tf.reduce_max(hough_Dx) * 3 / 5\n",
    "    hough_Dy_thresh = tf.reduce_max(hough_Dy) * 3 / 5\n",
    "\n",
    "    ax1.plot(hough_Dx.eval())\n",
    "    ax1.axhline(hough_Dx_thresh.eval(), lw=2, linestyle=':', color='r')\n",
    "    ax1.set_title('Hough Gradient X')\n",
    "    ax1.set_xlabel('Pixel')\n",
    "    ax1.set_xlim(0, a.shape[1])\n",
    "\n",
    "    ax2.plot(hough_Dy.eval())\n",
    "    ax2.axhline(hough_Dy_thresh.eval(), lw=2, linestyle=':', color='r')\n",
    "    ax2.set_title('Hough Gradient Y')\n",
    "    ax2.set_xlim(0, a.shape[0])\n",
    "    ax2.set_xlabel('Pixel')\n",
    "    # Get chess lines\n",
    "    lines_x, lines_y, is_match = getChessLines(hough_Dx.eval().flatten(),\n",
    "                                               hough_Dy.eval().flatten(),\n",
    "                                               hough_Dx_thresh.eval(),\n",
    "                                               hough_Dy_thresh.eval())\n",
    "\n",
    "    lines_x, lines_y, is_match = getChessLines(hough_Dx.eval().flatten(),\n",
    "                                               hough_Dy.eval().flatten(),\n",
    "                                               hough_Dx_thresh.eval() * .9,\n",
    "                                               hough_Dy_thresh.eval() * .9)\n",
    "\n",
    "    print (\"X\", lines_x, np.diff(lines_x))\n",
    "    print (\"Y\", lines_y, np.diff(lines_y))\n",
    "    if is_match:\n",
    "        print (\"Chessboard found\")\n",
    "    else:\n",
    "        print (\"Couldn't find Chessboard\")\n",
    "\n",
    "    # Plot blurred 1d hough arrays and skeletonized versions\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 5))\n",
    "\n",
    "    ax1.plot(hough_Dx.eval())\n",
    "    ax1.axhline(hough_Dx_thresh.eval(), lw=2, linestyle=':', color='r')\n",
    "    ax1.set_title('Hough Gradient X')\n",
    "    ax1.set_xlabel('Pixel')\n",
    "    ax1.set_xlim(0, a.shape[1])\n",
    "\n",
    "    ax2.plot(hough_Dy.eval())\n",
    "    ax2.axhline(hough_Dy_thresh.eval(), lw=2, linestyle=':', color='r')\n",
    "    ax2.set_title('Hough Gradient Y')\n",
    "    ax2.set_xlim(0, a.shape[0])\n",
    "    ax2.set_xlabel('Pixel')\n",
    "\n",
    "    # Plot lines for where peaks where found\n",
    "    if len(lines_x < 20):\n",
    "        for hx in lines_x:\n",
    "            ax1.axvline(hx, color='r')\n",
    "    if len(lines_y < 20):\n",
    "        for hy in lines_y:\n",
    "            ax2.axvline(hy, color='r')\n",
    "\n",
    "    plt.imshow(img)\n",
    "    for hx in lines_x:\n",
    "        plt.axvline(hx, color='b', lw=2)\n",
    "\n",
    "    for hy in lines_y:\n",
    "        plt.axhline(hy, color='r', lw=2)\n",
    "\n",
    "\n",
    "    print (\"X   (vertical)\", lines_x, np.diff(lines_x))\n",
    "    print (\"Y (horizontal)\", lines_y, np.diff(lines_y))\n",
    "    if is_match:\n",
    "        # Possibly check np.std(np.diff(lines_x)) for variance etc. as well/instead\n",
    "        print (\"7 horizontal and vertical lines found, slicing up squares\")\n",
    "        squares = getChessTiles(a, lines_x, lines_y)\n",
    "        print (\"Tiles generated: (%dx%d)*%d\" % (squares.shape[0], squares.shape[1], squares.shape[2]))\n",
    "    else:\n",
    "        print (\"Number of lines not equal to 7\")\n",
    "\n",
    "    letters = 'ABCDEFGH'\n",
    "\n",
    "    if is_match:\n",
    "        print (\"Order is row-wise from top left of image going right and down, so a8,b8....a7,b7,c7...h1\")\n",
    "        print (\"Showing 5 random squares...\")\n",
    "        for i in np.random.choice(np.arange(64), 5, replace=False):\n",
    "            print (\"#%d: %s%d\" % (i, letters[i % 8], i / 8 + 1))\n",
    "            display_array(squares[:, :, i], rng=[0, 255])\n",
    "    else:\n",
    "        print (\"Didn't have lines to slice image up.\")\n",
    "\n",
    "    img_save_dir = tiles_directory + \"/squares_%s\" % img_file[:-4]\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "    if not is_match:\n",
    "        print (\"No squares to save\")\n",
    "    else:\n",
    "        if not os.path.exists(img_save_dir):\n",
    "            os.makedirs(img_save_dir)\n",
    "            print (\"Created dir %s\" % img_save_dir)\n",
    "\n",
    "        for i in range(64):\n",
    "            sqr_filename = \"%s/%s_%s%d.png\" % (img_save_dir,\n",
    "                                               img_file[:-4], letters[i % 8], i / 8 + 1)\n",
    "            if i % 8 == 0:\n",
    "                print (\"#%d: saving %s...\" % (i, sqr_filename))\n",
    "\n",
    "            # Make resized 32x32 image from matrix and save\n",
    "            PIL.Image.fromarray(squares[:, :, i]) \\\n",
    "                .resize([32, 32], PIL.Image.ADAPTIVE) \\\n",
    "                .save(sqr_filename)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}