test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cv2\n",
    "import math\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import h5py\n",
    "import tensorflow as tf\n",
    "import imgaug as ia\n",
    "import imgaug.augmenters as iaa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from layers import dense_senet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "gpus = tf.config.experimental.list_physical_devices(device_type='GPU')\n",
    "cpus = tf.config.experimental.list_physical_devices(device_type='CPU')\n",
    "\n",
    "# tf.config.experimental.set_visible_devices(devices=cpus[0], device_type='CPU')\n",
    "# tf.config.experimental.set_visible_devices(devices=gpus[0], device_type='GPU')\n",
    "# for gpu in gpus:\n",
    "#     tf.config.experimental.set_memory_growth(device=gpu, enable=True)\n",
    "    \n",
    "tf.config.experimental.set_virtual_device_configuration(\n",
    "    gpus[0],\n",
    "    [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=6144)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "IMG_HEIGHT, IMG_WIDTH, IMG_CHANNEL = 160, 352, 1\n",
    "DATA_DIR = 'data/Shandong University/'\n",
    "CLASSES = 106 * 2 * 3\n",
    "IMG_PER_CLASS = 6\n",
    "DATA_LENGTH = CLASSES * IMG_PER_CLASS\n",
    "CACHE_FILE = 'data/cache.h5'\n",
    "TRAIN_RATIO = 0.8\n",
    "MODEL_FILE = 'data/model.h5'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prewitt算子/边缘检测\n",
    "def edge_detectio(grayImage):\n",
    "    kernelx = np.array([[1, 1, 1], [0, 0, 0], [-1, -1, -1]], dtype=int)\n",
    "    x = cv2.filter2D(grayImage, cv2.CV_16S, kernelx)\n",
    "    laplacian = cv2.Laplacian(grayImage,cv2.CV_64F)\n",
    "    absX= cv2.convertScaleAbs(laplacian)\n",
    "    absX= cv2.GaussianBlur(absX,(3,3),100)\n",
    "    h, w=absX.shape\n",
    "    return absX, h, w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#取边缘检测后边界的点\n",
    "def get_bound_piont(absX):\n",
    "    h, w=absX.shape\n",
    "    ret,absX=cv2.threshold(absX, 0, 255, cv2.THRESH_OTSU)\n",
    "    t, t1 =0, h-1\n",
    "    bound_up_x, bound_down_x, bound_up_y, bound_down_y = [], [], [], []\n",
    "    for i in range (w):\n",
    "        y_min, y_max = 0, h-1\n",
    "        for j in range (h):\n",
    "            if absX[j, i] > ret:\n",
    "                if j < h // 2 and j > y_min:\n",
    "                    y_min = j\n",
    "                if j > h // 2 and j < y_max:\n",
    "                    y_max = j\n",
    "                    break\n",
    "        if t < y_min:\n",
    "            t=y_min\n",
    "        if t1 > y_max:\n",
    "            t1=y_max\n",
    "        if y_max != h-1:\n",
    "            bound_up_x.append(i)\n",
    "            bound_up_y.append(y_max)\n",
    "        if y_min != 0:\n",
    "            bound_down_x.append(i)\n",
    "            bound_down_y.append(y_min)\n",
    "    t2=(t+t1)/2.\n",
    "    bound_up_x = np.array(bound_up_x)\n",
    "    bound_down_x = np.array(bound_down_x)\n",
    "    bound_up_y = np.array(bound_up_y)\n",
    "    bound_down_y = np.array(bound_down_y)\n",
    "    return bound_up_x, bound_down_x, bound_up_y, bound_down_y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fun2ploy(x,n):\n",
    "    '''\n",
    "    数据转化为[x^0,x^1,x^2,...x^n]\n",
    "    首列变1\n",
    "    '''\n",
    "    lens = len(x)\n",
    "    X = np.ones([1,lens])\n",
    "    for i in range(1,n):\n",
    "        X = np.vstack((X,np.power(x,i)))#按行堆叠\n",
    "    return X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def leastseq_byploy(x,y,ploy_dim):\n",
    "    '''\n",
    "    最小二乘求解\n",
    "    '''\n",
    "    #散点图\n",
    "    # plt.scatter(x,y,color=\"r\",marker='o',s = 50)\n",
    "\n",
    "    X = fun2ploy(x,ploy_dim)\n",
    "    #直接求解\n",
    "    Xt = X.transpose()#转置变成列向量\n",
    "    XXt=X.dot(Xt)#矩阵乘\n",
    "    XXtInv = np.linalg.inv(XXt)#求逆\n",
    "    XXtInvX = XXtInv.dot(X)\n",
    "    coef = XXtInvX.dot(y.T)\n",
    "\n",
    "    y_est = Xt.dot(coef)\n",
    "\n",
    "    return y_est,coef"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fitting_line(absX):\n",
    "    ploy_dim =2#拟合参数个数，即权重数量\n",
    "    ## 数据准备\n",
    "    bound_up_x, bound_down_x, bound_up_y, bound_down_y = get_bound_piont(absX)\n",
    "    x = bound_up_x\n",
    "    y = bound_up_y\n",
    "    x1 = bound_down_x\n",
    "    y1 = bound_down_y\n",
    "    # # 最小二乘拟合\n",
    "    [y_est,coef] = leastseq_byploy(x,y,ploy_dim)\n",
    "    [y_est1,coef1] = leastseq_byploy(x1,y1,ploy_dim)\n",
    "    # print(coef,coef1)\n",
    "    #显示拟合结果\n",
    "    # est_data = plt.plot(x,y_est-5,color=\"g\",linewidth= 2)\n",
    "    # est_data = plt.plot(x1,y_est1+5,color=\"g\",linewidth= 2)\n",
    "    b1,k1 = coef\n",
    "    b2,k2 = coef1\n",
    "    k = (k1+k2)/2.0\n",
    "    b = (b1+b2)/2.0\n",
    "    return k, b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_theta(w,k,b):\n",
    "    w_c = w/2\n",
    "    h_c = k * w_c + b\n",
    "    theta =math.atan(k)\n",
    "    # print(theta)\n",
    "    return theta,w_c,h_c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Rotating_picture(img,theta,h_c,w_c):\n",
    "    h, w,_=img.shape\n",
    "    matRotate = cv2.getRotationMatrix2D((h_c, w_c), 180*theta/math.pi, 1)\n",
    "    img1=cv2.warpAffine(img, matRotate, (w, h))\n",
    "    return img1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_point(w_c,h_c,x,y,angle):\n",
    "    (cX, cY) = (w_c, h_c)\n",
    "    if (cX-x)== 0:\n",
    "        theta1=math.atan((cY-y)/1)\n",
    "        l=1/math.cos(theta1)\n",
    "    else:\n",
    "        theta1=math.atan((cY-y)/(cX-x))\n",
    "        l=(cX-x)/math.cos(theta1)\n",
    "    # print(h,w)\n",
    "    new_x = cX-math.cos(angle +theta1)*l\n",
    "    new_y = cY-math.sin(angle +theta1)*l\n",
    "    # print(l, math.cos(theta1),cX,cY)\n",
    "    # print((new_x-cX)**2+(new_y-cY)**2,(x-cX)**2+(y-cY)**2)\n",
    "    return round(new_x), round(new_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Vertical_boundary(w_c,h_c,bound_down_x, bound_down_y,bound_up_x, bound_up_y,theta):\n",
    "    locs = list(zip(bound_down_x, bound_down_y))\n",
    "    locs1 = list(zip(bound_up_x, bound_up_y))\n",
    "    new_locs = list(map(lambda loc: get_point(w_c,h_c,loc[0], loc[1],   theta), locs))\n",
    "    new_locs1 = list(map(lambda loc: get_point(w_c,h_c,loc[0], loc[1],   theta), locs1))\n",
    "    new_locs = list(zip(*new_locs))\n",
    "    new_locs1 = list(zip(*new_locs1))\n",
    "    Y_max=int(max(new_locs[1]))\n",
    "    Y_min=int(min(new_locs1[1]))\n",
    "    return(Y_max, Y_min)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ncrease_contrast(img):\n",
    "    img = cv2.convertScaleAbs(img,alpha=1.5,beta=0)\n",
    "    return img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Horizontal_boundary(cropped):\n",
    "    max_1, max_2=0, 0\n",
    "    height, width,_ = cropped.shape\n",
    "    kernel_width = width // 20\n",
    "    x = [i for i in range(width)]\n",
    "    light = [0 for i in range(width)]\n",
    "    for w in range(width - kernel_width): #步长是1\n",
    "        light[w+kernel_width//2] = np.sum(cropped[:, w: w+kernel_width])\n",
    "        if w+kernel_width//2>width/3:\n",
    "            if light[w+kernel_width//2]>light[w+kernel_width//2+1] and light[w+kernel_width//2]>light[w+kernel_width//2-1]:\n",
    "               if max_1< light[w+kernel_width//2]:\n",
    "                    max_1=light[w+kernel_width//2]\n",
    "                    t800=w+kernel_width//2\n",
    "    cut_1, cut_2=int(t800/3), int(2/3*(width-t800)+t800)\n",
    "    return cut_1,cut_2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gray_normalization(img):\n",
    "    \"\"\" 灰度归一化 \"\"\"\n",
    "    g1, g2 = np.min(img), np.max(img)\n",
    "    img = ((img - g1) / (g2 - g1) * 255).astype(np.uint8)\n",
    "    return img\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def img_preprocess(img_path):\n",
    "    img = cv2.imread(img_path)\n",
    "    grayImage = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)\n",
    "    grayImage = cv2.GaussianBlur(grayImage,(3,3),0)\n",
    "    absX, h, w=edge_detectio(grayImage)\n",
    "    bound_up_x, bound_down_x, bound_up_y,       bound_down_y=get_bound_piont(absX)\n",
    "    k, b=fitting_line(absX)\n",
    "    theta,w_c,h_c=get_theta(w,k,b)\n",
    "    img=Rotating_picture(img,theta,h_c,w_c)\n",
    "    Y_max, Y_min=Vertical_boundary(w_c,h_c,bound_down_x,bound_down_y,bound_up_x,bound_up_y,-theta)\n",
    "    cropped = img[Y_max:Y_min, :]\n",
    "    X_min,X_max=Horizontal_boundary(cropped)\n",
    "    ncrease_contrast(img)\n",
    "    cropped = img[Y_max:Y_min, X_min:X_max]\n",
    "    # plt.imshow(cropped) # 裁剪坐标为[y0:y1, x0:x1]\n",
    "    cropped = cv2.resize(cropped, (IMG_WIDTH, IMG_HEIGHT))\n",
    "    # 对图片进行灰度归一化\n",
    "    cropped = gray_normalization(cropped)\n",
    "    cropped = cv2.cvtColor(cropped, cv2.COLOR_BGR2GRAY)\n",
    "    return cropped"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_path = 'data/Shandong University/001/left/index_1.bmp'\n",
    "img = img_preprocess(img_path)\n",
    "plt.imshow(img, cmap='gray')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(img.shape)\n",
    "print(img.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_data():\n",
    "    data = np.zeros(shape=(DATA_LENGTH, IMG_HEIGHT, IMG_WIDTH, IMG_CHANNEL), dtype=np.uint8)\n",
    "\n",
    "    for human_id in os.listdir(DATA_DIR):\n",
    "            human_dir = os.path.join(DATA_DIR, human_id)\n",
    "            for hand_id in os.listdir(human_dir):\n",
    "                hand_dir = os.path.join(human_dir, hand_id)\n",
    "                for finger_name in os.listdir(hand_dir):\n",
    "                    img_path = os.path.join(hand_dir, finger_name)\n",
    "                    img = img_preprocess(img_path)\n",
    "                    finger_type, finger_id = os.path.splitext(finger_name)[0].split('_')\n",
    "                    idx = (int(human_id) - 1) * 36 + \\\n",
    "                            {'left': 0, 'right': 1}[hand_id] * 18 + \\\n",
    "                            {'index': 0, 'middle': 1, 'ring': 2}[finger_type] * 6 + \\\n",
    "                            int(finger_id) - 1\n",
    "                    data[idx] = np.expand_dims(img, axis=-1)\n",
    "                    print('get data {}/{} '.format(idx+1, CLASSES*IMG_PER_CLASS), end='\\r')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(CACHE_FILE):\n",
    "    print(\"未发现处理好的数据文件，正在处理...\")\n",
    "\n",
    "    data = get_data()\n",
    "\n",
    "    h5f = h5py.File(CACHE_FILE, 'w')\n",
    "    h5f[\"X\"] = data\n",
    "    h5f.close()\n",
    "else:\n",
    "    h5f = h5py.File(CACHE_FILE, 'r')\n",
    "    data = h5f[\"X\"][:]\n",
    "    h5f.close()\n",
    "    print(\"发现处理好的数据文件，正在读取...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.dtype, data.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = data.reshape((CLASSES, IMG_PER_CLASS, IMG_HEIGHT, IMG_WIDTH, IMG_CHANNEL))\n",
    "permutation = np.random.permutation(CLASSES)\n",
    "train_data = data[permutation[: int(TRAIN_RATIO * CLASSES)]]\n",
    "test_data = data[permutation[int(TRAIN_RATIO * CLASSES): ]]\n",
    "train_size, test_size = train_data.shape[0], test_data.shape[0]\n",
    "print(train_data.shape, test_data.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(np.squeeze(data[1][1]), cmap='gray')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sometimes = lambda aug: iaa.Sometimes(0.5, aug)\n",
    "\n",
    "aug = iaa.Sequential(\n",
    "    [\n",
    "        sometimes(iaa.CropAndPad(\n",
    "            percent=(-0.05, 0.05),\n",
    "            pad_mode=ia.ALL,\n",
    "            pad_cval=(0, 255)\n",
    "        )),\n",
    "        sometimes(iaa.Affine(\n",
    "            scale={\"x\": (0.95, 1.05), \"y\": (0.95, 1.05)}, # scale images to 80-120% of their size, individually per axis\n",
    "            translate_percent={\"x\": (-0.05, 0.05), \"y\": (-0.05, 0.05)}, # translate by -20 to +20 percent (per axis)\n",
    "            rotate=(-3, 3), # rotate by -45 to +45 degrees\n",
    "            shear=(-3, 3), # shear by -16 to +16 degrees\n",
    "            order=[0, 1], # use nearest neighbour or bilinear interpolation (fast)\n",
    "            cval=(0, 255), # if mode is constant, use a cval between 0 and 255\n",
    "            mode=ia.ALL # use any of scikit-image's warping modes (see 2nd image from the top for examples)\n",
    "        )),\n",
    "        sometimes(iaa.OneOf([\n",
    "                    iaa.GaussianBlur((0, 3.0)), # blur images with a sigma between 0 and 3.0\n",
    "                    iaa.AverageBlur(k=(2, 7)), # blur image using local means with kernel sizes between 2 and 7\n",
    "                    iaa.MedianBlur(k=(3, 11)), # blur image using local medians with kernel sizes between 2 and 7\n",
    "        ])),\n",
    "        sometimes(iaa.Add((-10, 10), per_channel=0)),\n",
    "        # execute 0 to 5 of the following (less important) augmenters per image\n",
    "        # don't execute all of them, as that would often be way too strong\n",
    "#         iaa.SomeOf((0, 1),\n",
    "#             [\n",
    "#                 iaa.OneOf([\n",
    "#                     iaa.GaussianBlur((0, 3.0)), # blur images with a sigma between 0 and 3.0\n",
    "#                     iaa.AverageBlur(k=(2, 7)), # blur image using local means with kernel sizes between 2 and 7\n",
    "#                     iaa.MedianBlur(k=(3, 11)), # blur image using local medians with kernel sizes between 2 and 7\n",
    "#                 ]),\n",
    "# #                 iaa.Sharpen(alpha=(0, 1.0), lightness=(0.75, 1.5)), # sharpen images\n",
    "# #                 iaa.AdditiveGaussianNoise(loc=0, scale=(0.0, 0.02*255), per_channel=0), # add gaussian noise to images\n",
    "#                 iaa.Add((-10, 10), per_channel=0), # change brightness of images (by -10 to 10 of original value)\n",
    "# #                 iaa.AddToHueAndSaturation((-20, 20)), # change hue and saturation\n",
    "#             ],\n",
    "#             random_order=True\n",
    "#         )\n",
    "    ],\n",
    "    random_order=True\n",
    ")\n",
    "\n",
    "def amplify(X1, X2):\n",
    "    data_length = X1.shape[0]\n",
    "    flip_methods = [np.flipud, np.fliplr, np.flip, lambda _: _]\n",
    "    method_indices = np.random.choice(len(flip_methods), data_length)\n",
    "    for idx in range(data_length):\n",
    "        flip_method = flip_methods[method_indices[idx]]\n",
    "        X1[idx] = flip_method(X1[idx])\n",
    "        X2[idx] = flip_method(X2[idx])\n",
    "    return X1, X2\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# base_model = tf.keras.applications.DenseNet121(input_shape=(IMG_HEIGHT, IMG_WIDTH, 3), weights='imagenet', include_top=False)\n",
    "base_model = dense_senet.DenseNet121()\n",
    "# base_model.trainable = False\n",
    "input1 = tf.keras.Input(shape=(IMG_HEIGHT, IMG_WIDTH, IMG_CHANNEL))\n",
    "input2 = tf.keras.Input(shape=(IMG_HEIGHT, IMG_WIDTH, IMG_CHANNEL))\n",
    "\n",
    "x1, x2 = input1, input2\n",
    "\n",
    "# 1\n",
    "# x1 = base_model(x1)\n",
    "# x2 = base_model(x2)\n",
    "# x = tf.keras.layers.Concatenate(axis=-1)([x1, x2])\n",
    "# x = tf.keras.layers.GlobalAveragePooling2D()(x)\n",
    "# x = tf.keras.layers.Dense(512, activation='relu')(x)\n",
    "# # x = tf.keras.layers.Dense(512, activation='relu', kernel_regularizer=tf.keras.regularizers.l2(l=0.01))(x)\n",
    "# # x = tf.keras.layers.Dropout(0.3)(x)\n",
    "# # x = tf.keras.layers.Dense(128, activation='relu')(x)\n",
    "# x = tf.keras.layers.Dense(1, activation='sigmoid')(x)\n",
    "\n",
    "# 2\n",
    "x1 = base_model(x1)\n",
    "x2 = base_model(x2)\n",
    "x = tf.keras.layers.Subtract()([x1, x2])\n",
    "x = x**2\n",
    "x = tf.keras.layers.GlobalAveragePooling2D()(x)\n",
    "x = tf.keras.layers.Dense(128, activation='relu')(x)\n",
    "x = tf.keras.layers.Dense(1, activation='sigmoid')(x)\n",
    "\n",
    "# 3\n",
    "# x1 = base_model(x1)\n",
    "# x2 = base_model(x2)\n",
    "# x1 = tf.keras.layers.GlobalAveragePooling2D()(x1)\n",
    "# x2 = tf.keras.layers.GlobalAveragePooling2D()(x2)\n",
    "# # x1 = tf.keras.backend.l2_normalize(x1, axis=-1)\n",
    "# # x2 = tf.keras.backend.l2_normalize(x2, axis=-1)\n",
    "# x = tf.keras.layers.Subtract()([x1, x2])\n",
    "# x = x**2\n",
    "# x = tf.keras.layers.Dense(128, activation='relu')(x)\n",
    "# x = tf.keras.layers.Dense(1, activation='sigmoid')(x)\n",
    "\n",
    "outputs = x\n",
    "model = tf.keras.Model(inputs=[input1, input2], outputs=outputs)\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_epoch = 1000\n",
    "batch_size = 16\n",
    "learning_rate = 1e-5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = tf.keras.optimizers.Adam(learning_rate)\n",
    "loss_object = tf.keras.losses.BinaryCrossentropy()\n",
    "\n",
    "train_loss = tf.keras.metrics.Mean(name='train_loss')\n",
    "train_accuracy = tf.keras.metrics.BinaryAccuracy(name='train_accuracy', dtype=None, threshold=0.5)\n",
    "\n",
    "test_loss = tf.keras.metrics.Mean(name='test_loss')\n",
    "test_accuracy = tf.keras.metrics.BinaryAccuracy(name='test_accuracy', dtype=None, threshold=0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def batch_loader(data, batch_size=64):\n",
    "    data_length, classes, _, _, _ = data.shape\n",
    "    X1 = np.zeros(shape=(batch_size, IMG_HEIGHT, IMG_WIDTH, IMG_CHANNEL), dtype=np.float32)\n",
    "    X2 = np.zeros(shape=(batch_size, IMG_HEIGHT, IMG_WIDTH, IMG_CHANNEL), dtype=np.float32)\n",
    "    y = np.random.randint(2, size=(batch_size, )).astype(np.float32)\n",
    "    for idx, is_same in enumerate(y):\n",
    "        if is_same == 1.:\n",
    "            class_id = np.random.randint(data_length)\n",
    "            img_id1, img_id2 = np.random.choice(classes, 2, replace=False)\n",
    "            X1[idx] = data[class_id][img_id1]\n",
    "            X2[idx] = data[class_id][img_id2]\n",
    "        else:\n",
    "            class_id1, class_id2 = np.random.choice(data_length, 2, replace=False)\n",
    "            img_id1, img_id2 = np.random.randint(classes), np.random.randint(classes)\n",
    "            X1[idx] = data[class_id1][img_id1]\n",
    "            X2[idx] = data[class_id2][img_id2]\n",
    "    return X1, X2, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X1, X2, y = batch_loader(train_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# aug = iaa.AddToHueAndSaturation((-20, 20))\n",
    "\n",
    "X1, X2 = amplify(X1, X2)\n",
    "\n",
    "idx = 15\n",
    "plt.imshow(X1[idx].reshape(IMG_HEIGHT, IMG_WIDTH).astype(np.uint8), cmap='gray')\n",
    "plt.show()\n",
    "plt.imshow(X2[idx].reshape(IMG_HEIGHT, IMG_WIDTH).astype(np.uint8), cmap='gray')\n",
    "plt.show()\n",
    "print(y[idx])\n",
    "print(X1[idx].astype(np.uint8).shape)\n",
    "\n",
    "plt.imshow(aug(images=[X1[idx].astype(np.uint8)])[0].reshape(IMG_HEIGHT, IMG_WIDTH), cmap='gray')\n",
    "plt.show()\n",
    "plt.imshow(aug(images=[X2[idx].astype(np.uint8)])[0].reshape(IMG_HEIGHT, IMG_WIDTH), cmap='gray')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def train_on_batch(X1, X2, y):\n",
    "    with tf.GradientTape() as tape:\n",
    "        y_pred = model([X1, X2])\n",
    "        loss = loss_object(y_true=y, y_pred=y_pred)\n",
    "        loss = tf.reduce_mean(loss)\n",
    "    grads = tape.gradient(loss, model.variables)\n",
    "    optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))\n",
    "    \n",
    "    train_loss(loss)\n",
    "    train_accuracy(y, y_pred)\n",
    "    return loss\n",
    "\n",
    "@tf.function\n",
    "def test_on_batch(X1, X2, y):\n",
    "    y_pred = model([X1, X2])\n",
    "    t_loss = loss_object(y, y_pred)\n",
    "    \n",
    "    test_loss(t_loss)\n",
    "    test_accuracy(y, y_pred)\n",
    "    return t_loss\n",
    "\n",
    "for epoch in range(num_epoch):\n",
    "    \n",
    "    train_loss.reset_states()\n",
    "    train_accuracy.reset_states()\n",
    "    test_loss.reset_states()\n",
    "    test_accuracy.reset_states()\n",
    "    \n",
    "    for batch_index in range(train_size // batch_size):\n",
    "        X1, X2, y = batch_loader(train_data, batch_size=batch_size)\n",
    "        X1, X2 = amplify(X1, X2)\n",
    "        X1 = np.array(aug(images=X1.astype(np.uint8)), dtype=np.float32)\n",
    "        X2 = np.array(aug(images=X2.astype(np.uint8)), dtype=np.float32)\n",
    "        X1 /= 255.\n",
    "        X2 /= 255.\n",
    "        loss = train_on_batch(X1, X2, y)\n",
    "        template = '[Training] Epoch {}, Batch {}/{}, Loss: {}, Accuracy: {:.2%} '\n",
    "        print(template.format(epoch+1,\n",
    "                              batch_index,\n",
    "                              train_size // batch_size,\n",
    "                              loss,\n",
    "                              train_accuracy.result()),\n",
    "             end='\\r')\n",
    "\n",
    "    for batch_index in range(5*test_size // batch_size):\n",
    "        X1, X2, y = batch_loader(test_data, batch_size=batch_size)\n",
    "        X1 /= 255.\n",
    "        X2 /= 255.\n",
    "        loss = test_on_batch(X1, X2, y)\n",
    "        template = '[Testing] Epoch {}, Batch {}/{}, Loss: {}, Accuracy: {:.2%} '\n",
    "        print(template.format(epoch+1,\n",
    "                              batch_index,\n",
    "                              test_size // batch_size,\n",
    "                              loss,\n",
    "                              test_accuracy.result()),\n",
    "             end='\\r')\n",
    "        \n",
    "    template = 'Epoch {}, Loss: {}, Accuracy: {:.2%}, Test Loss: {}, Test Accuracy: {:.2%} '\n",
    "    print(template.format(epoch+1,\n",
    "                         train_loss.result(),\n",
    "                         train_accuracy.result(),\n",
    "                         test_loss.result(),\n",
    "                         test_accuracy.result()))\n",
    "    \n",
    "    model.save_weights(MODEL_FILE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "file_extension": ".py",
  "kernelspec": {
   "display_name": "Python 3.7.6 64-bit ('tf2': conda)",
   "language": "python",
   "name": "python37664bittf2condad2284dedc2eb448ea8250faba7f9d846"
  },
  "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.7.6"
  },
  "mimetype": "text/x-python",
  "name": "python",
  "npconvert_exporter": "python",
  "pygments_lexer": "ipython3",
  "version": 3
 },
 "nbformat": 4,
 "nbformat_minor": 4
}