{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cellular-edition",
   "metadata": {},
   "source": [
    "# QRBM digit generation demo\n",
    "\n",
    "Draw 0 or 1 on a drawing pad, click \"save\" and \"generate\" to see how RBM generates the number you drew."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "resistant-norwegian",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jupyter_drawing_pad as jd\n",
    "import ipywidgets as widgets\n",
    "from IPython.display import display,clear_output\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "\n",
    "widget = jd.CustomBox()\n",
    "btn_save = widgets.Button(description='save')\n",
    "btn_clear = widgets.Button(description='clear')\n",
    "btn_generate = widgets.Button(description='generate')\n",
    "output = widgets.Output()\n",
    "output2 = widgets.Output()\n",
    "box_layout = widgets.Layout(display='flex',\n",
    "                flex_flow='column',\n",
    "                align_items='center')\n",
    "box = widgets.HBox(children=[btn_save, btn_clear],layout=box_layout)\n",
    "vbox = widgets.VBox([widget.drawing_pad, box, output])\n",
    "\n",
    "PILimage = []\n",
    "dilation = np.zeros((28,28), dtype='int64')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "reserved-disease",
   "metadata": {},
   "outputs": [],
   "source": [
    "def on_button_save_clicked(b):\n",
    "    global PILimage, dilation, output\n",
    "    import numpy as np\n",
    "    import cv2 as cv\n",
    "    image = np.zeros((28, 28))\n",
    "    for x, y in zip(widget.drawing_pad.data[0], widget.drawing_pad.data[1]):\n",
    "        image[27-round(y*27/100),round(x*27/100)] = 1\n",
    "    kernel = np.ones((2,2),np.uint8)\n",
    "    dilation = cv.dilate(image,kernel,iterations = 1)\n",
    "#     dilation = dilation.astype('int64')\n",
    "    PILdilation = dilation.astype('uint8') * 255\n",
    "    PILimage = Image.fromarray(PILdilation)\n",
    "    PILimage = PILimage.resize((300,300))\n",
    "    with output:\n",
    "        clear_output()\n",
    "        display(PILimage)\n",
    "\n",
    "def on_button_clear_clicked(b):\n",
    "    global widget\n",
    "    widget.drawing_pad.clear()\n",
    "    \n",
    "def load_model(filename):\n",
    "    from qrbm.MSQRBM import MSQRBM\n",
    "    bm = MSQRBM(784, 30,qpu=False)\n",
    "    bm.load(f'./pretrained/{filename}')\n",
    "    return bm\n",
    "\n",
    "def on_button_generate_clicked(b):\n",
    "    global dilation, output2\n",
    "    image_height = 28\n",
    "    bm = load_model('01.txt')\n",
    "    generated_pic = bm.generate(test_img = dilation.flatten().tolist())\n",
    "#     plt.figure()\n",
    "#     plt.axis('off')\n",
    "    res = np.array(generated_pic, dtype='uint8') * 255\n",
    "    res = np.reshape(res, (28,28))\n",
    "    PILres = Image.fromarray(res)\n",
    "    PILres = PILres.resize((300,300))\n",
    "    with output2:\n",
    "        clear_output()\n",
    "        display(PILres)\n",
    "\n",
    "btn_save.on_click(on_button_save_clicked)\n",
    "btn_clear.on_click(on_button_clear_clicked)\n",
    "btn_generate.on_click(on_button_generate_clicked)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "metallic-market",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "vbox"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "exact-sending",
   "metadata": {},
   "outputs": [],
   "source": [
    "btn_generate"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "grave-judges",
   "metadata": {},
   "source": [
    "## Generated picture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "shaped-angle",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "output2"
   ]
  }
 ],
 "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.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}