diff --git a/mnist.ipynb b/mnist.ipynb
index 1315c1c..b929ff9 100644
--- a/mnist.ipynb
+++ b/mnist.ipynb
@@ -234,7 +234,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 34,
    "metadata": {},
    "outputs": [
     {
@@ -255,7 +255,19 @@
        "        font-size: 16px;\n",
        "        border-radius: 3px;\n",
        "        margin-top: 10px;\n",
-       "        width: 400px;\n",
+       "        width: 190px;\n",
+       "        margin-right: 20px;\n",
+       "    }\n",
+       "    #clear-button {\n",
+       "        background-color: #FF0000,; \n",
+       "        color: black;\n",
+       "        border: none;\n",
+       "        padding: 10px 20px;\n",
+       "        cursor: pointer;\n",
+       "        font-size: 16px;\n",
+       "        border-radius: 3px;\n",
+       "        margin-top: 10px;\n",
+       "        width: 190px;\n",
        "    }\n",
        "    #container {\n",
        "      display: flex;\n",
@@ -263,11 +275,20 @@
        "      align-items: center; /* Center horizontally */\n",
        "      justify-content: center;\n",
        "    }\n",
+       "    #btn-container {\n",
+       "      display: flex;\n",
+       "      flex-direction: row; /* Arrange children vertically */\n",
+       "      align-items: center; /* Center horizontally */\n",
+       "    }\n",
        "\n",
        "</style>\n",
        "<div id='container'>\n",
        "<canvas id=\"whiteboard\" width=\"400\" height=\"200\" fill_rect='white'></canvas>\n",
+       "<div id='btn-container'>\n",
        "<button id=\"capture-button\">Predict</button>\n",
+       "<button id=\"clear-button\">Clear</button>\n",
+       "</div>\n",
+       "\n",
        "</div>\n",
        "<script>\n",
        "    var canvas = document.getElementById('whiteboard');\n",
@@ -290,9 +311,19 @@
        "    canvas.addEventListener('mouseout', function () {\n",
        "        drawing = false;\n",
        "    });\n",
+       "    \n",
+       "    var clearButton = document.getElementById('clear-button');\n",
+       "    clearButton.addEventListener('click', function () {\n",
+       "       context.clearRect(0, 0, canvas.width, canvas.height);\n",
+       "    });\n",
+       "\n",
        "   var captureButton = document.getElementById('capture-button');\n",
        "    captureButton.addEventListener('click', function () {\n",
-       "      // need to image save option\n",
+       "       // Convert the canvas content to a data URL (image)\n",
+       "    var imageData = canvas.toDataURL(\"image/png\");\n",
+       "\n",
+       "    // Send the image data to the Jupyter kernel variable\n",
+       "    IPython.notebook.kernel.execute('image_data = \"' + imageData + '\"');\n",
        "    });\n",
        "</script>\n"
       ],
@@ -321,7 +352,19 @@
     "        font-size: 16px;\n",
     "        border-radius: 3px;\n",
     "        margin-top: 10px;\n",
-    "        width: 400px;\n",
+    "        width: 190px;\n",
+    "        margin-right: 20px;\n",
+    "    }\n",
+    "    #clear-button {\n",
+    "        background-color: #FF0000,; \n",
+    "        color: black;\n",
+    "        border: none;\n",
+    "        padding: 10px 20px;\n",
+    "        cursor: pointer;\n",
+    "        font-size: 16px;\n",
+    "        border-radius: 3px;\n",
+    "        margin-top: 10px;\n",
+    "        width: 190px;\n",
     "    }\n",
     "    #container {\n",
     "      display: flex;\n",
@@ -329,11 +372,20 @@
     "      align-items: center; /* Center horizontally */\n",
     "      justify-content: center;\n",
     "    }\n",
+    "    #btn-container {\n",
+    "      display: flex;\n",
+    "      flex-direction: row; /* Arrange children vertically */\n",
+    "      align-items: center; /* Center horizontally */\n",
+    "    }\n",
     "\n",
     "</style>\n",
     "<div id='container'>\n",
     "<canvas id=\"whiteboard\" width=\"400\" height=\"200\" fill_rect='white'></canvas>\n",
+    "<div id='btn-container'>\n",
     "<button id=\"capture-button\">Predict</button>\n",
+    "<button id=\"clear-button\">Clear</button>\n",
+    "</div>\n",
+    "\n",
     "</div>\n",
     "<script>\n",
     "    var canvas = document.getElementById('whiteboard');\n",
@@ -356,13 +408,61 @@
     "    canvas.addEventListener('mouseout', function () {\n",
     "        drawing = false;\n",
     "    });\n",
+    "    \n",
+    "    var clearButton = document.getElementById('clear-button');\n",
+    "    clearButton.addEventListener('click', function () {\n",
+    "       context.clearRect(0, 0, canvas.width, canvas.height);\n",
+    "    });\n",
+    "\n",
     "   var captureButton = document.getElementById('capture-button');\n",
     "    captureButton.addEventListener('click', function () {\n",
-    "      // need to image save option\n",
+    "       // Convert the canvas content to a data URL (image)\n",
+    "    var imageData = canvas.toDataURL(\"image/png\");\n",
+    "\n",
+    "    // Send the image data to the Jupyter kernel variable\n",
+    "    IPython.notebook.kernel.execute('image_data = \"' + imageData + '\"');\n",
     "    });\n",
     "</script>\n"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 36,
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "NameError",
+     "evalue": "name 'image_data' is not defined",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
+      "\u001b[1;32m/arun477/vae/mnist.ipynb Cell 13\u001b[0m in \u001b[0;36m<cell line: 9>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      <a href='vscode-notebook-cell://github/arun477/vae/mnist.ipynb#Y104sdnNjb2RlLXZmcw%3D%3D?line=5'>6</a>\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mtorch\u001b[39;00m\n\u001b[1;32m      <a href='vscode-notebook-cell://github/arun477/vae/mnist.ipynb#Y104sdnNjb2RlLXZmcw%3D%3D?line=7'>8</a>\u001b[0m \u001b[39m# Extract the base64 portion of the data URL\u001b[39;00m\n\u001b[0;32m----> <a href='vscode-notebook-cell://github/arun477/vae/mnist.ipynb#Y104sdnNjb2RlLXZmcw%3D%3D?line=8'>9</a>\u001b[0m image_data_base64 \u001b[39m=\u001b[39m re\u001b[39m.\u001b[39msub(\u001b[39m'\u001b[39m\u001b[39m^data:image/.+;base64,\u001b[39m\u001b[39m'\u001b[39m, \u001b[39m'\u001b[39m\u001b[39m'\u001b[39m, image_data)\n\u001b[1;32m     <a href='vscode-notebook-cell://github/arun477/vae/mnist.ipynb#Y104sdnNjb2RlLXZmcw%3D%3D?line=10'>11</a>\u001b[0m \u001b[39m# Decode the base64 string to bytes and create a PIL Image\u001b[39;00m\n\u001b[1;32m     <a href='vscode-notebook-cell://github/arun477/vae/mnist.ipynb#Y104sdnNjb2RlLXZmcw%3D%3D?line=11'>12</a>\u001b[0m image_bytes \u001b[39m=\u001b[39m base64\u001b[39m.\u001b[39mb64decode(image_data_base64)\n",
+      "\u001b[0;31mNameError\u001b[0m: name 'image_data' is not defined"
+     ]
+    }
+   ],
+   "source": [
+    "import numpy as np\n",
+    "from PIL import Image\n",
+    "import io\n",
+    "import re\n",
+    "import base64\n",
+    "import torch\n",
+    "\n",
+    "# Extract the base64 portion of the data URL\n",
+    "image_data_base64 = re.sub('^data:image/.+;base64,', '', image_data)\n",
+    "\n",
+    "# Decode the base64 string to bytes and create a PIL Image\n",
+    "image_bytes = base64.b64decode(image_data_base64)\n",
+    "image = Image.open(io.BytesIO(image_bytes))\n",
+    "\n",
+    "# Convert the PIL Image to a NumPy array\n",
+    "image_np = np.array(image)\n",
+    "image_tensor = torch.from_numpy(image_np)\n",
+    "image_tensor.shape"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,