grain_unearth-supervised_contrastive.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "executionInfo": {
     "elapsed": 2716,
     "status": "ok",
     "timestamp": 1624439925606,
     "user": {
      "displayName": "ramdhan wibawa",
      "photoUrl": "",
      "userId": "14519448541243845830"
     },
     "user_tz": -420
    },
    "id": "mW8mQghJfh6D"
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "import os\n",
    "from os import getenv\n",
    "from os.path import abspath, basename, split,dirname\n",
    "import tarfile\n",
    "from shutil import copyfile\n",
    "import random, glob\n",
    "from skimage.util import random_noise\n",
    "from skimage.transform import rotate\n",
    "import time\n",
    "import numpy as np\n",
    "import math\n",
    "import math as m\n",
    "\n",
    "import tensorflow as tf\n",
    "import tensorflow_datasets as tfds\n",
    "import tensorflow_addons as tfa\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras.layers import Layer\n",
    "\n",
    "from tensorflow.keras.applications import EfficientNetB0,EfficientNetB3, EfficientNetB2,EfficientNetB6\n",
    "from tensorflow.keras.applications.resnet50 import ResNet50\n",
    "from tensorflow.keras.applications.resnet import ResNet101\n",
    "from tensorflow.keras.applications.inception_v3 import InceptionV3\n",
    "from tensorflow.keras.applications.mobilenet import MobileNet\n",
    "from tensorflow.keras.applications.densenet import DenseNet121, DenseNet169, DenseNet201\n",
    "from tensorflow.keras import layers\n",
    "from tensorflow import keras\n",
    "#from tensorflow.keras.applications.efficientnet import preprocess_input\n",
    "#from tensorflow.keras.applications.resnet50 import preprocess_input\n",
    "from tensorflow.keras.applications.densenet import preprocess_input\n",
    "from keras.preprocessing.image import ImageDataGenerator\n",
    "import tensorflow.keras.backend as K\n",
    "from sklearn.metrics import accuracy_score\n",
    "from tensorflow.keras.callbacks import ModelCheckpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_cosine_schedule_with_warmup(lr, num_warmup_steps, num_training_steps, num_cycles=0.5):\n",
    "    \"\"\"\n",
    "    Modified the get_cosine_schedule_with_warmup from huggingface for tensorflow\n",
    "    (https://huggingface.co/transformers/_modules/transformers/optimization.html#get_cosine_schedule_with_warmup)\n",
    "\n",
    "    Create a schedule with a learning rate that decreases following the\n",
    "    values of the cosine function between 0 and `pi * cycles` after a warmup\n",
    "    period during which it increases linearly between 0 and 1.\n",
    "    \"\"\"\n",
    "\n",
    "    def lrfn(epoch):\n",
    "        if epoch < num_warmup_steps:\n",
    "            return float(epoch) / float(max(1, num_warmup_steps)) * lr\n",
    "        progress = float(epoch - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))\n",
    "        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))) * lr\n",
    "\n",
    "    return tf.keras.callbacks.LearningRateScheduler(lrfn, verbose=True)\n",
    "\n",
    "def acc_grain(y_true, y_pred):\n",
    "    y_true_class = K.argmax(y_true, axis=-1)\n",
    "    y_pred_class = K.argmax(y_pred, axis=-1)\n",
    "    sample_weight = np.hstack((np.where(y_true_class.numpy().reshape(-1,1)==7,True,False),np.where(y_true_class.numpy().reshape(-1,1)==24,True,False),np.where(y_true_class.numpy().reshape(-1,1)==13,True,False)))\n",
    "    sample_weight = np.any(sample_weight,axis=1)\n",
    "    sample_weight = np.where(sample_weight,8.94,1.0)\n",
    "    return K.constant(accuracy_score(y_true_class,y_pred_class,sample_weight=sample_weight))\n",
    "\n",
    "def outer_product(x):\n",
    "    #Einstein Notation  [batch,1,1,depth] x [batch,1,1,depth] -> [batch,depth,depth]\n",
    "    phi_I = tf.einsum('ijkm,ijkn->imn',x[0],x[1])\n",
    "    \n",
    "    # Reshape from [batch_size,depth,depth] to [batch_size, depth*depth]\n",
    "    phi_I = tf.reshape(phi_I,[-1,x[0].shape[3]*x[1].shape[3]])\n",
    "    \n",
    "    # Divide by feature map size [sizexsize]\n",
    "    size1 = int(x[1].shape[1])\n",
    "    size2 = int(x[1].shape[2])\n",
    "    phi_I = tf.divide(phi_I, size1*size2)\n",
    "    \n",
    "    # Take signed square root of phi_I\n",
    "    y_ssqrt = tf.multiply(tf.sign(phi_I),tf.sqrt(tf.abs(phi_I)+1e-12))\n",
    "    \n",
    "    # Apply l2 normalization\n",
    "    z_l2 = tf.nn.l2_normalize(y_ssqrt, axis=1)\n",
    "    return z_l2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 19182 files belonging to 28 classes.\n",
      "Found 7648 files belonging to 28 classes.\n"
     ]
    }
   ],
   "source": [
    "num_classes = 28\n",
    "image_size = 224  # We'll resize input images to this size\n",
    "input_shape = (224,224, 3)\n",
    "learning_rate = 0.001\n",
    "batch_size = 64\n",
    "hidden_units = 512\n",
    "projection_units = 128\n",
    "num_epochs = 10\n",
    "dropout_rate = 0.5\n",
    "temperature = 0.05\n",
    "\n",
    "size = (image_size, image_size)\n",
    "\n",
    "train_ds = tf.keras.preprocessing.image_dataset_from_directory(\n",
    "    \"D:/Ramdhan/Unearth/visual-grain-analysis/data/public/train\",\n",
    "    image_size=size,\n",
    "    batch_size=batch_size,\n",
    "    interpolation='nearest'\n",
    ")\n",
    "\n",
    "val_ds = tf.keras.preprocessing.image_dataset_from_directory(\n",
    "    \"D:/Ramdhan/Unearth/visual-grain-analysis/data/public/val\",\n",
    "    image_size=size,\n",
    "    batch_size=batch_size,\n",
    "    interpolation='nearest'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rotate_tensor(image):\n",
    "    degree = tf.random.normal([])*360\n",
    "    image = tfa.image.rotate(image, degree * math.pi / 180)\n",
    "    return image\n",
    "\n",
    "def input_preprocess_train(image, label):\n",
    "    image = rotate_tensor(image)\n",
    "    image = tf.cast(image,'float32') #activate for densenet\n",
    "    image = preprocess_input(image)\n",
    "    return image, label\n",
    "\n",
    "def input_preprocess_val(image, label):\n",
    "    image = tf.cast(image,'float32')\n",
    "    image = preprocess_input(image)\n",
    "    return image, label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds_c = train_ds.map(input_preprocess_train)\n",
    "val_ds_c = val_ds.map(input_preprocess_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(64, 224, 224, 3) (64,)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(-0.5, 223.5, 223.5, -0.5)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAOcAAADnCAYAAADl9EEgAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4xLjMsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+AADFEAAAgAElEQVR4nO1d69arKAwNZ837P/LH/GitEHKFgNi6Z535WgWCyCY3tCnnDA8ePNgP/67uwIMHD2g85HzwYFM85HzwYFM85HzwYFM85HzwYFP8J53MGZ5QLoGUgte0BJAGhzq9/3KtJOa42uDx4R3V9/ayKp9tvTjqfK4p/zml3gsp0bfn0ZwepH/xxKTEOM4npbx23taB/Dnmbastr9Mb93nFmO+I37zqDqT0DxL4NYcJRaMJ/a36gM6XZTK0fSOJJDXCdq4u5CZ71dYDKx5yGnCs3Mc0vcMUY4l5MPK4EIXRCRiby4vMfnnA4CGngtKkkrTamIxzsmb0F8vuEwAFMYuDVKMNbw4Gn+W91GpEd+AXTdvfu2Ij0iL/spFLfLfMa2ytVsinBmzOKSvOi5q1vdDDs5aglNpW2vgxgv7W1RqxehJIU7RLS+Fj6Tx6xlyNjm6gndAS1N/2LxH0d67UiGu0ZTtJE9i1ptAwrUblAvKKMMrVxve0Lz+/5qk+5CxwCTETwDHtQn1ZIvVxSmmSFWUBdLzwOdsa/m51Vix98V/Rnr9xlQZcdcNxwj0SSWiZ1tUF8ud/xYFEWsPufonaXMevEDRJz3P+wg6hqJvc6haL8HpnUBhBE/WRbr2KE1N3m4ogde4WImVXjfhH4Bt2Dz07hAhErr4zNJ9LEGGtasQ0t92Uk3SyHRFj9s0a9HuvTMHlN/Xta7rMWmlnD6HG7J5sZtug856z8PWGmgs/Sc7LiQnwmYfD0dhBmHQgQ9AQKrH+pz2Su8X9nIDvvCoB1huZwTM9ejpSt2ySo6Y4/MkX8/URW5ZCzNLP/6ojYLmW6h59IUG/74oEeIj5qTOlJ7l80MMuJzUf0Nf8afvYFeSCVjwXhQJXre6ob3Pgu6bzd12NAI2YlKacF+QZaZkjJjpmENHs6N3B5Uu+TnwzQb/nSgRYiNnUmdOVs32PgGbjup64se7HdVfKdcolapxStaD4dw1Ve4u/hKBfm+fc+galTG5wl+tQJXP7sUql2KmTuTynWi8Wde5zjPp3yYH+VJ5za2ICgHsrHOdr4jJNjtNKnX6KRVoYLScHX92y/TyQce/eE9j+hjj3mb+2t3JOJL9h3oNR7Rdm2pJHf5eg9+05ge1vxHvjAX6Zlzi51bnZ7gPybAjM5Ydl25x40EbCbxL0nr1GiHgwerpz/XnAoyO9YWDvsZmgfpOCdlVFfDrBgkGwwWLF/wJuT86IVZF7NUgoMrzSBMmhNYWsCe1TZrTFVmo99jG1aTwKaviO2vN+PS4QoS2XKovsmGuM+qgV6esbHV7SJKX6+gMGIpKg7aah8Q7ejaD36m2BGWbsaivKJo+elOUGgnMDBZNmIevPWZamERQAfo2g9+npG7P8yyXERELYqUZ0Ro9k+rKa4yVi63GYcV/uQtB79PKN25JS6oBYjo4G0cEeq+eMdCZj21NPpFm6P4VMXT25P25DzihiUntn1xC0nVC9G/C4M7KpSkzqSXP8qvH04A7ac/8eQjuQ3rWTeNppvX+ZnPtpi6vEnmSj+fivTQ1vAOio7R2zO2RBdifo1r3j/EvrRCmnZZ1iuB56P+jeVikUxpIlqNh+cw7E5YZk1y59Gbs/B7ptz6JWtZ0IaUahYinfr7vZgTpXj2H1UE4HMvHvbHxPGmzZq0hi9iB0UwKxOthTKIKfKnQumkg97fWYwiZ0borXHy7Yjwrb9ehqP6AkZsjEyuWmOm8v2iCPlp8093mynVru8523C8l+EVrJqy0DCluRcxdict+74H0KBRVgtxW0e/nE5q5KQMyf9NpmC74EjkNcPf8wtujNVb/oVWLexPWmUKgCNkObapessZil4eLY8eGPlr6zZG5fPQ9LXN6T8odpr1jdp8ok3nig16E+ptqcJTptJuYFmLKlz9AoDmZpVT7DuglBr+1F+tfMs5EgjqZfKLNVWlGHJ5W288YkgCBmn2hX/Z1B74OnL8xKyKb2BgT97zLJ6OJHiJCJz/omtvMvlh230ie+J8Kpsg/qBvV0lGsXFuP8vSVSAqhff+W/a7sPx/LlIaV/AOlfpalWRcoorXhplE5YFcjVPKNyQud3I2bkOJ/mZ39datdYg4u151Lph385mn/UjjVyDcfZBHUvpAeryzyDq01oO43bxt+vCtMiRHaBTnXqEqSobdVudfI6gi6THLlx3bNJoBx4XE9qJ2Iyubaue1csRWvm48NGmBIYqsBHbTVi8rXhMoIukeohJkU+L3koc9lKTNPNMoDVZIlxNvNZzmTqMxeQmw97QHGx+9okB8kmRTBCmArrCTpdopeY+LNVO0b4ryF+KNOJjD8xF2aOXBNyNuNjg2j/vvd6KXfGU2cVppLT+4teXTI662ntdBM9Q7Nd79OWsUHzwoRCtJ3bTm8Lz5hWdaCdczhYSAYPF2vPadJ6ftHL1T7Y78vSuSppTYl1ghYkrQIhGLQ0BH4xPL4nrkdtUtDSaisJOiXP2fPDQea2neWXmyw5Nw1GmNrUCXzuRxRmhU9KhckdZZDfrcSdE2Ma6d+S32EJXwZ6iSn5e94gjSclskzBUEt1CUNncZFfJKMP7l8oBYA6nczOowUaNExzWtS9lZilX4BTIdJgeybrHFK2Paj8QGfIslmU8CYEqdlNmRsdtU2AdwodR32gxpGbb6uGNoScPcS0kE2z/70w1++dQYSp6UJu2yBTI4ikJEFn5C4CsGGX1JTdVUM5rJspYkrml2bd9UIbPJcsLnc4Aq0D74FRiQlUAULEjiyYBW2FFwJkVjeJbnuuaTvUeknM0s/j/EZPdNWTXrFsSJgOy8+lO4qYNxOg+FNqPuyFGd1iUypCsGI0fXeO8zyCdpm11KsqP+eo8sZ2zQl4ojyFiImgmd4mec6OeHb5sK/b3FRzzupWAnj9gFv5K26GjR5eGWT99A9gQvTWTXtKWwKMEUFKCEt1KChWTHefNIwsSl0CpSY2JSbAPIXe7MBy9sHqaq00cc2ak9OWpEllhJeQR50oU5WT719wYkxaVzkAgHSfKG0JVgOFtanbO1xAUuqT2t9gDaqSUwr4jGrLjwxHnZmklI5JrVGmZeX7zSJMPs3aG3CyQTRJ6w0JNoJWdSMQSFBRF3MpklHN1UNMb1lOHpY9fVLPdLJS/fWuCO17x3iH36IgE/eyJ0lnR1FxxJdz6W47qXP155aas0T8fdBbDHDzWckRe3Bj99YaVFGUacq1zaHsWpx8xaRdhLsTcw7omRa5qWD2fR6ndxXmqi95pabiZOFND6HyBxLbkbI/H3+cpTQV22+zifnRyIPac3ATwpEAbg2EFeaWJGO+9uKvbIXmJFNGt7XRXwjpftNI3AzEY45jGaSZPEDQ7prpoymL7uQ5syNDe/EcMb3ase/WbaqiNu3WKmg5XjwvPJtLpL0ephRLB/p8zgSQIRdR67brEZsAON+gN9pL1S+/30bxpLbP5w6Za7oUgdGu65d/+qFc2dSU5Fvydc6fYvGT89N7/29nSQjZJrewD5DosqsILm4VvDFBR1DnOQ+022QswUHq+HAwyUnQPn0bdPM5UxVAX9lmYEXOdRZy8+F+GB1TIiYpyuhxf1YOr5+cg71jHWeQVzGLLym1rSGCbFdqzSSdvAlCJr4xIGQZppH5xMLhfy7dhKDlIaljnkHsgXer3q1V0y8g0KqLAE7lvT7YaLeMnNzFRm4FtBzHssu/1hrcY1oraMs+IgY2s25nhFgvTSP2VqM0JWXpta6wTr3p5NQudmQgegg/KzI7ezviR0hEmY0xulB751MkIcu/lJy2kky/KeS0XLDVj5RkcO1KGCMRvV1P6k80NDnp87/fA2058CMWTUg87KPE7yYnJdjakTvPnfZNby+s0pqmbWM3Nm2Hu268EV452Hfs0pRkwzwFu8hJJfE5c9ESYaXa9Mi3yupD3UPJ55uKXP3hTt+amBGwRvR722zbP+fH67+4W+DahGAyqYAnKE7yU1vxsONMmQqS7FigHm6g8qV004/zUgE3++gRPcaTv+Xlb47PiWSImhObrBQ8vqM1LYK/c7a7P9rqk91IuMKkTeTHBwxot+Ok2TmG9Ky1mKz0TMTtj0M1ayU1be0IVX80LxlFTMkkP49eaNIaF+Vv0JpzNiF0FWFRm641wbm52Suv+2HrGblJq8w4jsiGSyXZ8l7auT14YAGdszBV5YM7+r3HBnM5V6dHa6lIFYZVw3om2mjKRcLxC1QJHaVLSu30yrfVLzW2OsY3ZvGcTQgc6JGswzrHZypW286a4x/l2fZAfsEX0Q3KJzz+Rs2LmYQ8wPeXPsqlUPgavNzyr1afW5FF9A7c7EFXEGGb8HGDulRqStcOXG26Sm2dZzSueOFKpZRR1/JSvMTUgksRkAbG19883DEqKq0ivUUWsSitC5/zI0m8CxGiOdUS9eCk90z2LNSWUhH5fpPPSaU4ohA9J8pJzE3oHpmaWctpOI2QZLsHIdN5vifSzaK8mRuQssToPMtNI3xr6V0wAX42WbereoOcHktIJad1tbdqo0jzd76MIvPPNIgPe7SjREz2PAMPMRMA/0aZDcK+QwR1rVCvvzZT2B9rkeparE6RnN4BsizGMzQlbjtahhRk6LEqxP69G/Rew6cPUkeQJs74c3kxkoqYiFFxKUnxgXZUqaKa5WXt4+j6F/veWpgVVZUHyeayy23xrdtKWglqkR06hqnmG+5jdQyvbhdo0VnuUwmJkNKxUVJ62gAIfCrFG1OwdlKKDh9yPTLdE99Yofd6GlmdzCS1puCzYlG7uJ4riEmNhaWeBVZiUtFdjGHN2XNTLSSJirZaUxYtjmBBHMS2klZGVmcjmpusu4Hv2QMp5VXOm1WEpNqxzqnlO4Q0YmkX0iu3q16KCy7FmLPKdEHO0i7a0IOZZu2UAJtRruSadQWEJGEWeHwxrOE8fmYT2BD66PU5b0PMsg1np8kxvdDvnIOIEZbrSfMWiHOaPPuP5xrLcT6i1I5GTGtOKaPPlE0vX0chOUjtmJpx+LUztGET+FiR72IwJ1rLX0zk2tM7bzm4t+9J0ARz7WA5ntUlDmho86L5qZifnrxpD8jF9MZak/Q5mQHOxKcRlJZbRIsh0doMMZ3RtK4HfmLFLwlqH5IsjYpMR421JDdJJzcH+fY94loyLlOdsbkQONodvZAOR2tnRLSunxcJ0uAjYlG+am+kT4Nkgl0//nPBX58tQae5ZlHoJmevCWtpZ0SHjU9e/g174SgYQo5DvAt8th3c3r7Ixf8BuLBNmThrfPCm3Bq4ySl1bCR3iev3DEDMhIsZfnNMpQxpI1+vZ6NFL27sZlbg347I//AWl/scDeiMwuxzapa4FtSQBgbb754B8CSVbQ2OD/9xPeV1s2PA5YTwOQNGen7070qNGjHx1Tfis5+kvkR6+rVcabyHNr5bbqS2WnllUm3ETaiYHUFUYItsF6srtIlArc8050HZ9tWm7lzt/bpSS2ScQ1TI0Dre7mitpeEM/FojpWYswY95E+ntm0yyW8hm4xdkM7CmpIKcd4T3/tkirKmi9aosgkrOhP5pAno6Ts1RStbcifRKOF720mjqQsGXZx4RCXDpWrEIrbH6+ivHqqM1Zi2bh2jWeieF1wy2pgnEfoTN3Hy901VgdTeae3chS2dGRNt2y4jcIb3uC1/XDi8xAQKf55R8SymnVpYLFbwZrNe3ipRcHOrq4ZwpP7NXrdWLku9D2A4hCtpqsUsg4ujVLJNWbRYFf46/0TtOJPE7EBNg3lyQickTNROfqNZ1+X4MaU5vamVPYu4FPC6rCJPZL/eB/PrS/CaoN39Qn6ddszmzeOgnAClIOUuSkFuwdJ5wq9Ysxw1Hu6MDN2L+9WJizoyE1jG3Mv5K18RzWYuZRKNLc1o6hf0YkYPCgrYmRpOnCPGYs+9eiPWjJsMx3GQQaBf7tgdq37kNem0zZw2upfkI3b6HoSpFxXZbM0fmELOjFxUk62NkweK8rarxCzHSBe0VJXgTAmXVie33dGoA5pdKa8AXZ548mfz4aWPFgMzgpllrEkyzuAVegpon4mpnNxCe12JaYrVXD8FQQGh4Umc9rbKBUnPD1GdmdlhX9h5ims3kC2flqC/NwTt+dD4U15zrdJl+n5NCRJfw5Ju5G0bvwUVAzp+UchqBOZm+gVnbDaPWPIvm5hgdbOOWt7mz1bzxnd9DUZdPxF+pXam9FYjObXY1JwSGRrqXmc/U9++FlLvkHyPTZ/F8qNv31EjrG3jF5y7NM/GitUeLGK+2u29MRapXI2OhupGbMXVmrKFdsCjy2UYbz+/oPps2vmuwhptLzWtpd34KdH9ijuY0xYlD+RUboPd6aSvooGAuyFjVYiTaBmPmkE1/h5AQlBTbXGJQbDIZAWJXX8mcJYVy43BB5HYoIKSaBtiuO0BdqDz7zD78AJa9Q0jyU7m6c4k5/mD1cPAisV+7ZHIxAlefsKDNTN4etNv2LHFrGzFnwkzOUptp8GjI3rouUPeiUy3P0ujs2BbCPOTlxtbMtQtIOUNkbjSk5Gy1sKxRR1ipXgTGZ4jrNSURaRWN5KETnw0V76cOLFkAaWy0wMTMIEsURvtIbUJIBVlyJUUnEtWfutbYiGrz3GXW9viPJagcXg/c8gNnZejiEdjYcDO7M9cAfXcQXsLkUZO1ZXvkRdy4QFLXb6WMEDMCZvmC0MteR4JhtIKwVsQBCU9aynbiGszQ8GXK5CCQlubj+lb+HemsJVljemTMGoQYKWORtQOu6t8npZLb49bgG9noZujtkqQ10/v/mFwyMakk1pmQOV+hOi8HNfQOIS4wjc9b8RBTQW4/dvdpQ2IOgdRiPWYrfTQVmhcgn5M+ZTh/LzEW7oetKb8xgph3wOXXdCzS1KJewNzPy1ebFnGLjYeYcqgzYfVTKktiwYyC+ZGxcmGKJOaG8+MaaI4jM+jdMaVNw7c9XcoAxPW0I0ObrgAk09goS0ZNf3QpKXME5lQKFVLuxbWE7J+R0/tdzhVDN7n+mKpvSMwhCEniTKZObLFYKir78TdzmUyJJSbA5J+dH60zC1GR2rh0c9GQYrKWRSmYOPdlxOSpdiY3ckMebjuJRMzShp1LTACHz5mKfxK4BLh6cAl7x7QmZUmEdltrtLgBdCzx3ojt/3nHzs0HVJn2O+/aH4M/MnmNqy8YNadHtFi2XKwkR3ZDWP3tMAiWl8Rd4KsZTl6LXjdYe3fQ8fN/5+sxlfbINuqVc4VzpJIzZAJS1gTA1hNFwnRSAtBOftZNcnZIbzrWJpCsPvTlqQ1tdh9dpty4sAqiWRtOTIDbT5LpaSPK6qGy5rP7cQG6I7UsLDNYy9bTslJVb87Ih/wcQwNu60XwNfiaGxdOrZxLCJFP+V0yb8La7kAjmeO0hMxKQsr7herWpfM0Wa0xmxJhP2RUYdIMvnqOSdY4t+5GxvFwW9j/Zc2uTXOaYfhcH2eS0l4kVZKmUB0oaR9DqzsTFVKZozkngArE+Cb93HDwEnMX6JutatSejniX+QCMm7XU6FBLGEU+TXpNaGlookIqczSnEx7t0kfMQ8o84JUyFEZfkyVpeYIqxFVcqG27RTGBoFZfcqkULf5+ji6XGZ2FrTQnd4Mic4ozHxWbPpcz1It8IZA1tLi5RymR8hhn4U3CuJi2s3r3ucFpB7hfKfTjcnJqsbLoOSLlxGYgpN+U1VUMmBgQ7+3ADUxaABBWIrYQc54iZKr+rsYlZq2kIfH6FTksVzxgPc0MIhrFsUePX8qeX7SYdbsEufkgtExpR6wTE/pWD/Q2eU4v8NrDRTUplNYUY7kNY7XWDIPgClHmrJeYgM/jm7AI3kVM796x6R3X6rPH6ODTPPSRk1m18V88KbzZoQsChuEY7j/XADrOxSeP71TMkhVx0cC7pztb4XUBmVyiNDM1V01ftE4BQA85i+vSSIc1oOd+x86Na1TmNH9TkIV1AveZqsc2eiu8Ol3vok2fa9f9x/N8Jj5R5cbA31ybz0lEHLxmJ+dTUFF+Db7o7V1tWQHvwaR8aE4jcmOljs60/NA42ldhvjpb+94jJPLNNCrVYjG9OejkJJbWnnuFDQx8Edrlc/6rSxvfTRMIOcmIazHxblNiAtDEfH2yJFBK4lkCRi0oxeIjpgzZrA0iJtWkZ271BsivnFlhvmZPZKeApah4TxYtaGN3yrbMtx4CJib2MO0aU5bz8oC9kDVnICkB+u5zPzHLUudauWKuhfmazDGP1tTMWS7C2zSymQattaYWfaWe46SIWX6m2+qx8OonWOw3z/yCrx70+JOjMqWWZhNzR6uZ84Hcfd2QoBJaDcmV4Egq19ZlHnWPO+BngUrOUWLizxZoMu2Xum42LSNmst/mqKDdbMTILGOrHv1Wkqcd2R5NWfciVd88cP2QkYaICWohZq+szEQ4R9HbpHstdQqiyCbpEBGTWNsVXGTyPrlajrCDzpm1rRrxLHzSmdEfmQzZvjc/49MjB7U2SbWNEPP4a2ojkR/V9pkmxHK2k7sB/wonJiUVWWsjtdgFwK5ZPSQU8ftMWApD5FxBSj8hmRqTtOaoP22u+xYklacnEC/ne4jJdVdzrvggEDZ8bTKkY350k3NvYq6ZWUuDXOlsg5NrIeZunIvxc49nLWX6ZHaOtPq2rFG3cpSSosMxcL19rzvSBz4N42mfHr75oY0oYprGxRj21tpyjcii6FAEMY+WUuF30hsReuwKKli0Jvzn+pWxGf4VHi4fMct6csuRO2qiiNmVzhDa59ozEyDGVZqKNhiEfUb6AnJVtoQnzWLFiBo7Ie4QyvmvEdcDab3xEjMX/856kqEXg/I2ecZCKmtqR9GUkkZux0lpstTQt9CaHjkSMTP6d5isFMm0O0ppXLpH2kCrT6WU1SMIqsmy9ONo7/XPtjFq9FlOXH10Ypk9FiXonKBtS5qGqrV6A+3ZgteWPDG1tji/En/n2CFpT5tG1gNC+Q8gnRyOvHcmUhHHfPJfAxEZqbU2NWzOKgWl09wUEiGFe4PRI8KzwNIhHuriyos2OvdN3caJMtTVy9ie5yzM26NLo+glZre8wMYsVl8oMXP7kY410lXNl76ImLFo0x9+YtbteInJhSVHYX/YuiDoSGBImizJWAa3ByDNqTmzjUvaWPuuAl8YQ1DN51QLYtzKpG3JQadT+CitFD7SZZc+KtWGZSnl4XsTQv4bjtiOlOGsf6v/FkVTr/FT1nGBY35xnLr2TH0ZtcUDESfiNRi5WSotq5GFmPizRnwuCOTVyi9Me/setuBLUMc1i4qbgB4zMUIp9AzzLGVELRLkGA4P3u7AF4RnVvm3rpXYI9SoYruNIiaWK0Mq5X6HUEb+Z3UO9LVGi11R6E+UELbgILzrqhuaGVCcx0aVSsxyHl1EzN4x4mIGvK93xPLly639U+tyl1E5iZj8zNBuQZfmzPkPUqp5za0fUmcoLST5UGU5FzEBwvfVUiTtWzuJhh0VTTEcSVOq7I5FZMwpf1rEMmgt2X6jNF9JZe1u9jg4fGtNmSyEMTO3WfGo/CZoVKwBGwoA7fD5EEtQPqzQSu0WZVndPANiXS0WEHNEVDtN06cty7yiB8LTE84E0aD3MMEfWWDI58woB6qW/3Sm7XL5XdKwXQjQmhZihsxvqRFqYdcIV5alyi8kZYy4JHzDsqQ7ZrXvqJmqSZbas2P8je+CD8pWQX8BJrtAebx9y6Sym9yOxriOeIh5tGdzhaYhRpze8fz+7wWP82S17dY47Jf+kFEEYeTbdEpYNQ/DiMmtXJb5k4nPAPSqiOVuiNOkbcM75yVlZj5ge81z8VyQqK4xa9hiyNmhPVdoMimWF4UppKeIabWfOZ/A7pxth1esANvn73Ofb9aQoksy8Tk3Jehg0zjiNKdAULzeiXMEpQlIUW1xvXuTJt7UQCdFSEoTgnCs9DXLz5tqSR6tGXAe8ZgEUtS1PVrr6HaFa4npHFiBN7GbEN6CcJqlhLp4Z74cHnq+LfomzXjBV38kWYAlaMNZaVoWgIpsLUiljDRNLaxtBBaDi4TVZq62T0iyvZRkBg+jpbnFz84foFIp5bkSWnB6pY0WroQsXddCx9Qx7vym5iyARMwD3MrCLZvlz8drxJTyCxQMM8HhAk4JCL1SLEoZ4h9n9trnTtlai51eizkVUhDISurbANvqtAOFj3BzjCc+1sTWUGQpxBebmbe39u8PEtC7iChomQODROBX0Tkolc4285szOy7Umr0ifHECbHfxxOT6xEd6qe+8Nm2GtSNgCjA5lXLsw+XWmTYwjuoLddqy6+mxFTFT9Yc818yljYk5hpaYJWzRVclMLonZzu4IYgIsyHNKG+XFeuh7OSSvf/T7vaVwZHTEdpugZ0HMLHVKCgZtCMoNqdMm5WfaYMUWPk2jjEpIThaWK5QZICbAhQEhi7akM1t16Tbexs+22/0+pxXvKDQ5RtI1TzZrR5vmfnahtVmk2XS8Y4qL2FJtc4R3OOyDxAS4aIeQl5i+VqjW5mCbICdFTEqDSi7ULYCJKV0AJqNETNoG4+ULcvNfCDEBVpHT+A4iyzqFa8jB8LnYUhFbXaiForvaIleb8pe7ZIO1nUW4DEXI2lelMpvtmaLNIFIeWKc5CYKORvUfYgJtmZXQo/6hkDI2oy2fl0JpT04TcuBNYj7MQw3u+1gwMQFWm7XEBVgWeiqq20PMGVv4trIML+xMqVPWJrLaUKG+dMrn5UfNiGMTiAlwhc+JLsQzjFwcll4n25KRASHrNJiKkgnUICwiK2c4rpN6oLznXLzfGqvgBjLX/yYRE+DiR8YkyMOY2TmZizKNYTIhlXIp5tmQl4K7T/R4UyuUrSYhGeiZx2jNWU9TvHENOY8N8paigGl2vgqRyllpsbZvm8gN7NG0EKwcTtoyomL85efUnKHD1h47KE3VmAeu2/hueMUJdeO1nUDqEAebtpdrTwA2eL4J2eQAAAWmSURBVDgTq9e4Vh6XqyxxnsskiTlic6mV9/EFxAS42Kyldg954mzc0ZXZhMsV8QXEvAZUVPXUinzqo/xbHreGItHxyaZsict9TupnBrXL32UOWsIL04Ht/ss7NAZ57pd5yNZcrdMsPQ65QNiclxITYJPnOaMvecX83IoDC7XnNQtj/bPyeOtJbj5pdpczcLSYlAcu15wAEGjDU4EiLCtG0i7aezXWL0qnf/iyVGr3pSVmWYdry4FF/iWFPcgJsG4QtlJ5AbBkEoIwO49ZK6g6GltG6FtSUj4lKUE4RxW/jpgAO5ETIGQwVnJvbdJd6ET5dyJmj227SaSMllL7dkYvmgklBW5eH8Fe5AQQB0XfTYRznUT87ts2IizEbP63eWvspuBtehjCFjvxeLHCbUDKA/uRkwG/Tbk8nohbNo8+2xBzckeWWQiNWXtoywy1/W5qALXFfS+080WBHw57klPYf8utl/Uejy2SHOuw15wKQk9KxLLtDrc5d/P6CPYkJwCUW/wkQwaa42tJ+ZW8QFgxovorMDHwAoxtK2lJJ3zMDbFFnpNF/oPs+BUzehLNo08mPn+jvr5iAaJlOv1Hsez6TQVe7Ks5DxCrmt3tXzP4FEkf2KFzhHj8r/pkNX0LO2xzYgLcgZwAFUEpo4XXVtfosf1vux3XX0tLSnp/kJWYcAtiAtyFnG9o7r2v1hx8Uyhq1RSmHoJPwkiev7/ZsW/2JsQE2N3nLIEeMbNs1roCZdr8gQ7+9Zdk6f4F40akPHArzZnzX7NePkSYg2VTOdFfzvxmwBaoGxIT4GbkBACA/KduSEAV2iNf+eRGDJb3PddfMuCXQPvudtv+fe/G/cgJLw2KbxV9C/gb8xC0xcJtui85jaDjmRPsHHTuUboxMQHu5HMSwJu6eHO3fS3Fip9mePxPH+rfw3wdkcFs6bs5KQ/cUnMCQPULZuR5gM8TgHT9Kd2qcCdirp7O/ENe5Z4wbyvpa4gJcGNyAoBz29Xje26FTH6E1v6xPPLw/rfpNrxe3JucAOIN4faUVNUfgl6uNeWSkidc3NMvIybAzX3OD5jXbFpSLqt8T60fV2GImB1Oda6EJuRnUj5n/bnyRr+QkCXurzkPmG7U9dTYSZMOE7OnSmXOSo8OYObXryn5dmICfBM5AdhN8ji1XVVZzJbrl4cABI1ZOxYcMbH87ycmwLeRE6C5cbn5RNz0hQTdRXOO9eNd27vSZOmr9ChDqW5/g5gA30hOAOIGYmLSzzWswtUEHTNn+2rz4Rwq3UUEgjZ56dZKfCc5AeB8kwLeBgbATs8f0KC9cvP7jecf2ji0JvY1y0WSeAUbnPfr6mXsWnwvOaEkZnsG44rcdYZaR1Dno+V11cvlbtdEUsqOk5jntZdaktgd9EUbCzz4anLm4qa2E4Fes69EBp2wy/DuRDmGH1J2c5OqWBKSCAj9KDEBvpycAPB5zOzzXdi3mejDS7DjFCxTHb3EbLfl4bPUBvf3bp8fJibAt2xC0FBsUrhbKkOIMXe140EqHYMOFzBXH07TtSZsLfFV5LcCPxy+XnN+0Dxmxpi1m7J3qQ4phH2IWf7tbKv+ypi4DzE/+B1yNqCn+86W1BJfVM46uWB/S+Ihe+PBvwA/Rc6X/6lPgN3niLd75vLHhQcQU1pIyGjvozEb/IbPWSLn92536wO9e4JKElFbyO0NVpteY8z7j59Q9gb9ns1DSha/R054RWzPbdRsoW39TwrUFnJbRSLBFEbMA3VE9kPX3U2Ui/FTZu0HKL1CFlnSkTGgDW6dDVC/zBYNInv7EFPFT2rOE/ffIjbSe+x/f3zBQauh5l3bWH5MWRN+U3MCmHydr17cpWsbDARxjeX89xDTgZS/egY+eHBf/K7mfPBgczzkfPBgUzzkfPBgUzzkfPBgUzzkfPBgUzzkfPBgU/wP2Lu6/v5iExwAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 288x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "for x_train, y_train in train_ds_c.take(1):\n",
    "    print(x_train.shape, y_train.shape)\n",
    "    \n",
    "plt.figure(figsize=(4, 4))\n",
    "image = x_train[np.random.choice(range(x_train.shape[0]))]\n",
    "plt.imshow(image.numpy().astype(\"uint8\"))\n",
    "plt.axis(\"off\")"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "model1 = ResNet50(weights='imagenet', include_top=False, input_shape=(224, 224, 3))\n",
    "    \n",
    "for layer in model1.layers:\n",
    "    print(layer.name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ArcFace(Layer):\n",
    "    '''Custom Keras layer implementing ArcFace including:\n",
    "    1. Generation of embeddings\n",
    "    2. Loss function\n",
    "    3. Accuracy function\n",
    "    '''\n",
    "\n",
    "    def __init__(self, output_dim, class_num, margin=0.5, scale=64., **kwargs):\n",
    "        self.output_dim = output_dim\n",
    "        self.class_num = class_num\n",
    "        self.margin = margin\n",
    "        self.s = scale\n",
    "\n",
    "        self.cos_m = tf.math.cos(margin)\n",
    "        self.sin_m = tf.math.sin(margin)\n",
    "        self.mm = self.sin_m * margin\n",
    "        self.threshold = tf.math.cos(tf.constant(m.pi) - margin)\n",
    "        super(ArcFace, self).__init__(**kwargs)\n",
    "\n",
    "\n",
    "    def build(self, input_shape):\n",
    "        # Create a trainable weight variable for this layer.\n",
    "        self.kernel = self.add_weight(name='kernel', \n",
    "                                      shape=(input_shape[1], self.class_num),\n",
    "                                      initializer='glorot_normal',\n",
    "                                      trainable=True)\n",
    "        super(ArcFace, self).build(input_shape)  # Be sure to call this at the end\n",
    "\n",
    "\n",
    "    def call(self, x):\n",
    "        embeddings = tf.nn.l2_normalize(x, axis=1, name='normed_embeddings')\n",
    "        weights = tf.nn.l2_normalize(self.kernel, axis=0, name='normed_weights')\n",
    "        cos_t = tf.matmul(embeddings, weights, name='cos_t')\n",
    "        return cos_t\n",
    "\n",
    "\n",
    "    def get_logits(self, labels, y_pred):\n",
    "        cos_t = y_pred\n",
    "        cos_t2 = tf.square(cos_t, name='cos_2')\n",
    "        sin_t2 = tf.subtract(1., cos_t2, name='sin_2')\n",
    "        sin_t = tf.sqrt(sin_t2, name='sin_t')\n",
    "        cos_mt = self.s * tf.subtract(tf.multiply(cos_t, self.cos_m), tf.multiply(sin_t, self.sin_m), name='cos_mt')\n",
    "        cond_v = cos_t - self.threshold\n",
    "        cond = tf.cast(tf.nn.relu(cond_v, name='if_else'), dtype=tf.bool)\n",
    "        keep_val = self.s*(cos_t - self.mm)\n",
    "        cos_mt_temp = tf.where(cond, cos_mt, keep_val)\n",
    "        mask = tf.one_hot(labels, depth=self.class_num, name='one_hot_mask')\n",
    "        inv_mask = tf.subtract(1., mask, name='inverse_mask')\n",
    "        s_cos_t = tf.multiply(self.s, cos_t, name='scalar_cos_t')\n",
    "        output = tf.add(tf.multiply(s_cos_t, inv_mask), tf.multiply(cos_mt_temp, mask), name='arcface_logits')\n",
    "        return output\n",
    "\n",
    "\n",
    "    def loss(self, y_true, y_pred):\n",
    "        #labels = K.argmax(y_true, axis=-1)\n",
    "        #logits = self.get_logits(labels, y_pred)\n",
    "        #loss = tf.losses.categorical_crossentropy(labels=labels, logits=logits)\n",
    "        loss = tf.losses.categorical_crossentropy(y_true, y_pred,from_logits=True)\n",
    "        return loss\n",
    "\n",
    "\n",
    "    def accuracy(self, y_true, y_pred):\n",
    "        labels = K.argmax(y_true, axis=-1)\n",
    "        logits = self.get_logits(labels, y_pred)\n",
    "        accuracy = categorical_accuracy(y_true=labels, y_pred=logits)\n",
    "        return accuracy\n",
    "    \n",
    "\n",
    "    def compute_output_shape(self, input_shape):\n",
    "        return (input_shape[0], self.output_dim)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_encoder(): \n",
    "    inputs = keras.Input(shape=input_shape)\n",
    "    resnet = DenseNet121(include_top=False, weights='imagenet', input_tensor=inputs, pooling=\"avg\")\n",
    "    for layer in resnet.layers:\n",
    "        layer.trainable = True\n",
    "    #augmented = data_augmentation(inputs)\n",
    "    outputs = resnet.output\n",
    "    model = keras.Model(inputs=inputs, outputs=outputs, name=\"grains-encoder\")\n",
    "    return model\n",
    "\n",
    "def create_encoder_bilinear():\n",
    "    input_tensor = layers.Input(shape=input_shape)\n",
    "    model1 = DenseNet121(weights='imagenet', include_top=False, input_tensor=input_tensor,input_shape=input_shape)\n",
    "    model2 = DenseNet121(weights='imagenet', include_top=False, input_tensor=input_tensor,input_shape=input_shape)\n",
    "\n",
    "    for layer in model1.layers[:-4]:\n",
    "        layer.trainable = False\n",
    "    for layer in model2.layers[:-7]:\n",
    "        layer.trainable = False\n",
    "        \n",
    "    for i, layer in enumerate(model1.layers):\n",
    "        layer._name = 'model1_' + layer.name\n",
    "\n",
    "    last_layer1 = model1.get_layer('model1_relu')\n",
    "    last_output1 = last_layer1.output\n",
    "\n",
    "    for i, layer in enumerate(model2.layers):\n",
    "        layer._name = 'model2_' + layer.name\n",
    "\n",
    "    last_layer2 = model2.get_layer('model2_relu')\n",
    "    #last_layer2 = model2.get_layer('model2_top_conv')\n",
    "    last_output2 = last_layer2.output\n",
    "    \n",
    "    \n",
    "    model1_ = tf.keras.Model(inputs=model1.input, outputs=last_output1)\n",
    "    model2_ = tf.keras.Model(inputs=model2.input, outputs=last_output2)\n",
    "    \n",
    "    #optimizer = tf.keras.optimizers.Adam(learning_rate=1e-5, decay=1e-3)\n",
    "   \n",
    "    \n",
    "    model1_.compile(optimizer='adam',loss=\"categorical_crossentropy\", metrics=[acc_grain],run_eagerly=True)\n",
    "    model2_.compile(optimizer='adam',loss=\"categorical_crossentropy\", metrics=[acc_grain],run_eagerly=True)\n",
    "    \n",
    "    d1=model1_.output\n",
    "    d2=model2_.output\n",
    "\n",
    "    bilinear = layers.Lambda(outer_product, name='outer_product1')([d1,d2])\n",
    "    \n",
    "    model = tf.keras.Model(inputs=model1.input, outputs=bilinear, name=\"grains-encoder\")\n",
    "    \n",
    "    return model\n",
    "\n",
    "def create_encoder_MDNet(img_dim=input_shape, num_pipelines=4, decrease_by=0.25):\n",
    "\n",
    "    assert img_dim[0] == img_dim[1]\n",
    "    assert num_pipelines > 1\n",
    "    pipelines = []\n",
    "\n",
    "    model_input = layers.Input(shape = img_dim)\n",
    "    \n",
    "    model = DenseNet121(weights='imagenet', include_top=False, input_shape=(None, None, 3), pooling='avg')\n",
    "    \n",
    "    for layer in model.layers:\n",
    "        layer.trainable = True\n",
    "\n",
    "    # Number of pipelines\n",
    "    for pipeline_idx in range(num_pipelines):\n",
    "\n",
    "        new_dim = int(img_dim[0] * (pipeline_idx * decrease_by)) if pipeline_idx != 0 else img_dim[0]\n",
    "        crop_by = int(new_dim/2) if pipeline_idx != 0 else 0\n",
    "\n",
    "        x = layers.Cropping2D(cropping = crop_by, name = 'Cropped_by_' + str(crop_by))(model_input)\n",
    "        \n",
    "        out = model(x)\n",
    "        \n",
    "        pipelines.append(out)\n",
    "    \n",
    "    model_output = layers.Concatenate(axis = -1)(pipelines)\n",
    "    \n",
    "    model = tf.keras.Model(inputs = model_input, outputs = model_output, name =\"grains-encoder\")\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "encoder = create_encoder()\n",
    "classifier = create_classifier(encoder)\n",
    "classifier.summary()\n",
    "\n",
    "history = classifier.fit(train_ds, batch_size=batch_size, epochs=num_epochs)\n",
    "\n",
    "accuracy = classifier.evaluate(val_ds)[1]\n",
    "print(f\"Test accuracy: {round(accuracy * 100, 2)}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SupervisedContrastiveLoss(keras.losses.Loss):\n",
    "    def __init__(self, temperature=1, name=None):\n",
    "        super(SupervisedContrastiveLoss, self).__init__(name=name)\n",
    "        self.temperature = temperature\n",
    "\n",
    "    def __call__(self, labels, feature_vectors, sample_weight=None):\n",
    "        # Normalize feature vectors\n",
    "        feature_vectors_normalized = tf.math.l2_normalize(feature_vectors, axis=1)\n",
    "        # Compute logits\n",
    "        logits = tf.divide(\n",
    "            tf.matmul(\n",
    "                feature_vectors_normalized, tf.transpose(feature_vectors_normalized)\n",
    "            ),\n",
    "            self.temperature,\n",
    "        )\n",
    "        return tfa.losses.npairs_loss(tf.squeeze(labels), logits)\n",
    "    \n",
    "class triplet_loss(keras.losses.Loss):\n",
    "    def __init__(self, margin=1.0, kind='hard', name=None):\n",
    "        super(triplet_loss, self).__init__(name=name)\n",
    "        self.margin = margin\n",
    "        self.kind = kind       \n",
    "\n",
    "    def __call__(self, labels, feature_vectors, sample_weight=None):\n",
    "        # Normalize feature vectors\n",
    "        feature_vectors_normalized = tf.math.l2_normalize(feature_vectors, axis=1)\n",
    "        \n",
    "        if self.kind == 'hard':\n",
    "            loss = tfa.losses.triplet_hard_loss(labels, feature_vectors_normalized, margin=self.margin, soft=False)\n",
    "        elif self.kind == 'soft':\n",
    "            loss = tfa.losses.triplet_hard_loss(labels, feature_vectors_normalized, margin=self.margin, soft=True)\n",
    "        elif self.kind == 'semihard':\n",
    "            loss = tfa.losses.triplet_semihard_loss(labels, feature_vectors_normalized, margin=self.margin)\n",
    "            \n",
    "        return loss\n",
    "    \n",
    "def add_projection_head(encoder):\n",
    "    inputs = keras.Input(shape=input_shape)\n",
    "    features = encoder(inputs)\n",
    "    outputs = layers.Dense(projection_units, activation=\"relu\")(features)\n",
    "    model = keras.Model(inputs=inputs, outputs=outputs, name=\"encoder_with_projection-head\")\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-23-17bb7203622b>:1: is_gpu_available (from tensorflow.python.framework.test_util) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.config.list_physical_devices('GPU')` instead.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.test.is_gpu_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"encoder_with_projection-head\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "input_2 (InputLayer)         [(None, 224, 224, 3)]     0         \n",
      "_________________________________________________________________\n",
      "grains-encoder (Functional)  (None, 1024)              7037504   \n",
      "_________________________________________________________________\n",
      "dense (Dense)                (None, 128)               131200    \n",
      "=================================================================\n",
      "Total params: 7,168,704\n",
      "Trainable params: 7,085,056\n",
      "Non-trainable params: 83,648\n",
      "_________________________________________________________________\n",
      "\n",
      "Epoch 00001: LearningRateScheduler reducing learning rate to 0.0.\n",
      "Epoch 1/10\n",
      "300/300 [==============================] - 244s 812ms/step - loss: 4.5754 - val_loss: 4.7972\n",
      "\n",
      "Epoch 00002: LearningRateScheduler reducing learning rate to 2.5e-05.\n",
      "Epoch 2/10\n",
      "277/300 [==========================>...] - ETA: 15s - loss: 3.0038"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-24-66b9f9bae906>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m     14\u001b[0m history = encoder_with_projection_head.fit(train_ds_c,validation_data=val_ds_c, batch_size=batch_size, \n\u001b[0;32m     15\u001b[0m                                            \u001b[0mcallbacks\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mlr_schedule\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;31m#steps_per_epoch=19182 / batch_size,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 16\u001b[1;33m                                            epochs=10)\n\u001b[0m\u001b[0;32m     17\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     18\u001b[0m \u001b[0mprint\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m'duration run'\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mtime\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mtime\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m-\u001b[0m\u001b[0mstart\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m/\u001b[0m\u001b[1;36m60\u001b[0m\u001b[1;33m/\u001b[0m\u001b[1;36m60\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;34m' hours'\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python36\\site-packages\\tensorflow\\python\\keras\\engine\\training.py\u001b[0m in \u001b[0;36m_method_wrapper\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m    106\u001b[0m   \u001b[1;32mdef\u001b[0m \u001b[0m_method_wrapper\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    107\u001b[0m     \u001b[1;32mif\u001b[0m \u001b[1;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_in_multi_worker_mode\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m  \u001b[1;31m# pylint: disable=protected-access\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 108\u001b[1;33m       \u001b[1;32mreturn\u001b[0m \u001b[0mmethod\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    109\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    110\u001b[0m     \u001b[1;31m# Running inside `run_distribute_coordinator` already.\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python36\\site-packages\\tensorflow\\python\\keras\\engine\\training.py\u001b[0m in \u001b[0;36mfit\u001b[1;34m(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_batch_size, validation_freq, max_queue_size, workers, use_multiprocessing)\u001b[0m\n\u001b[0;32m   1101\u001b[0m               \u001b[0mlogs\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtmp_logs\u001b[0m  \u001b[1;31m# No error, now safe to assign to logs.\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1102\u001b[0m               \u001b[0mend_step\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mstep\u001b[0m \u001b[1;33m+\u001b[0m \u001b[0mdata_handler\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mstep_increment\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 1103\u001b[1;33m               \u001b[0mcallbacks\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mon_train_batch_end\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mend_step\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mlogs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   1104\u001b[0m         \u001b[0mepoch_logs\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mcopy\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcopy\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mlogs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1105\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python36\\site-packages\\tensorflow\\python\\keras\\callbacks.py\u001b[0m in \u001b[0;36mon_train_batch_end\u001b[1;34m(self, batch, logs)\u001b[0m\n\u001b[0;32m    438\u001b[0m     \"\"\"\n\u001b[0;32m    439\u001b[0m     \u001b[1;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_should_call_train_batch_hooks\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 440\u001b[1;33m       \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_call_batch_hook\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mModeKeys\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mTRAIN\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m'end'\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mlogs\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mlogs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    441\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    442\u001b[0m   \u001b[1;32mdef\u001b[0m \u001b[0mon_test_batch_begin\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mlogs\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;32mNone\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python36\\site-packages\\tensorflow\\python\\keras\\callbacks.py\u001b[0m in \u001b[0;36m_call_batch_hook\u001b[1;34m(self, mode, hook, batch, logs)\u001b[0m\n\u001b[0;32m    287\u001b[0m       \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_call_batch_begin_hook\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mmode\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mlogs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    288\u001b[0m     \u001b[1;32melif\u001b[0m \u001b[0mhook\u001b[0m \u001b[1;33m==\u001b[0m \u001b[1;34m'end'\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 289\u001b[1;33m       \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_call_batch_end_hook\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mmode\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mlogs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    290\u001b[0m     \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    291\u001b[0m       \u001b[1;32mraise\u001b[0m \u001b[0mValueError\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m'Unrecognized hook: {}'\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mformat\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mhook\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python36\\site-packages\\tensorflow\\python\\keras\\callbacks.py\u001b[0m in \u001b[0;36m_call_batch_end_hook\u001b[1;34m(self, mode, batch, logs)\u001b[0m\n\u001b[0;32m    307\u001b[0m       \u001b[0mbatch_time\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtime\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mtime\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;33m-\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_batch_start_time\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    308\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 309\u001b[1;33m     \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_call_batch_hook_helper\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mhook_name\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mlogs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    310\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    311\u001b[0m     \u001b[1;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_check_timing\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python36\\site-packages\\tensorflow\\python\\keras\\callbacks.py\u001b[0m in \u001b[0;36m_call_batch_hook_helper\u001b[1;34m(self, hook_name, batch, logs)\u001b[0m\n\u001b[0;32m    343\u001b[0m       \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    344\u001b[0m         \u001b[1;32mif\u001b[0m \u001b[0mnumpy_logs\u001b[0m \u001b[1;32mis\u001b[0m \u001b[1;32mNone\u001b[0m\u001b[1;33m:\u001b[0m  \u001b[1;31m# Only convert once.\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 345\u001b[1;33m           \u001b[0mnumpy_logs\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtf_utils\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mto_numpy_or_python_type\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mlogs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    346\u001b[0m         \u001b[0mhook\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mbatch\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mnumpy_logs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    347\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python36\\site-packages\\tensorflow\\python\\keras\\utils\\tf_utils.py\u001b[0m in \u001b[0;36mto_numpy_or_python_type\u001b[1;34m(tensors)\u001b[0m\n\u001b[0;32m    535\u001b[0m     \u001b[1;32mreturn\u001b[0m \u001b[0mt\u001b[0m  \u001b[1;31m# Don't turn ragged or sparse tensors to NumPy.\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    536\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 537\u001b[1;33m   \u001b[1;32mreturn\u001b[0m \u001b[0mnest\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmap_structure\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0m_to_single_numpy_or_python_type\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtensors\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    538\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    539\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python36\\site-packages\\tensorflow\\python\\util\\nest.py\u001b[0m in \u001b[0;36mmap_structure\u001b[1;34m(func, *structure, **kwargs)\u001b[0m\n\u001b[0;32m    633\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    634\u001b[0m   return pack_sequence_as(\n\u001b[1;32m--> 635\u001b[1;33m       \u001b[0mstructure\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m[\u001b[0m\u001b[0mfunc\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0mx\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;32mfor\u001b[0m \u001b[0mx\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mentries\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    636\u001b[0m       expand_composites=expand_composites)\n\u001b[0;32m    637\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python36\\site-packages\\tensorflow\\python\\util\\nest.py\u001b[0m in \u001b[0;36m<listcomp>\u001b[1;34m(.0)\u001b[0m\n\u001b[0;32m    633\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    634\u001b[0m   return pack_sequence_as(\n\u001b[1;32m--> 635\u001b[1;33m       \u001b[0mstructure\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m[\u001b[0m\u001b[0mfunc\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0mx\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;32mfor\u001b[0m \u001b[0mx\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mentries\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    636\u001b[0m       expand_composites=expand_composites)\n\u001b[0;32m    637\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python36\\site-packages\\tensorflow\\python\\keras\\utils\\tf_utils.py\u001b[0m in \u001b[0;36m_to_single_numpy_or_python_type\u001b[1;34m(t)\u001b[0m\n\u001b[0;32m    531\u001b[0m   \u001b[1;32mdef\u001b[0m \u001b[0m_to_single_numpy_or_python_type\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mt\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    532\u001b[0m     \u001b[1;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mt\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mops\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mTensor\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 533\u001b[1;33m       \u001b[0mx\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mt\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mnumpy\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    534\u001b[0m       \u001b[1;32mreturn\u001b[0m \u001b[0mx\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mitem\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;32mif\u001b[0m \u001b[0mnp\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mndim\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mx\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;33m==\u001b[0m \u001b[1;36m0\u001b[0m \u001b[1;32melse\u001b[0m \u001b[0mx\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    535\u001b[0m     \u001b[1;32mreturn\u001b[0m \u001b[0mt\u001b[0m  \u001b[1;31m# Don't turn ragged or sparse tensors to NumPy.\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python36\\site-packages\\tensorflow\\python\\framework\\ops.py\u001b[0m in \u001b[0;36mnumpy\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m   1061\u001b[0m     \"\"\"\n\u001b[0;32m   1062\u001b[0m     \u001b[1;31m# TODO(slebedev): Consider avoiding a copy for non-CPU or remote tensors.\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 1063\u001b[1;33m     \u001b[0mmaybe_arr\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_numpy\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m  \u001b[1;31m# pylint: disable=protected-access\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   1064\u001b[0m     \u001b[1;32mreturn\u001b[0m \u001b[0mmaybe_arr\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcopy\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mmaybe_arr\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mnp\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mndarray\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;32melse\u001b[0m \u001b[0mmaybe_arr\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1065\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python36\\site-packages\\tensorflow\\python\\framework\\ops.py\u001b[0m in \u001b[0;36m_numpy\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m   1027\u001b[0m   \u001b[1;32mdef\u001b[0m \u001b[0m_numpy\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1028\u001b[0m     \u001b[1;32mtry\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 1029\u001b[1;33m       \u001b[1;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_numpy_internal\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   1030\u001b[0m     \u001b[1;32mexcept\u001b[0m \u001b[0mcore\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_NotOkStatusException\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[1;33m:\u001b[0m  \u001b[1;31m# pylint: disable=protected-access\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1031\u001b[0m       \u001b[0msix\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mraise_from\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mcore\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_status_to_exception\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0me\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcode\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0me\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmessage\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;32mNone\u001b[0m\u001b[1;33m)\u001b[0m  \u001b[1;31m# pylint: disable=protected-access\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "encoder = create_encoder()\n",
    "#encoder = create_encoder_bilinear()\n",
    "#encoder =  create_encoder_MDNet()\n",
    "\n",
    "encoder_with_projection_head = add_projection_head(encoder)\n",
    "encoder_with_projection_head.compile(optimizer='adam',\n",
    "    loss=SupervisedContrastiveLoss(temperature),)\n",
    "    #loss= triplet_loss(margin=1.0, kind='semihard'),)\n",
    "\n",
    "lr_schedule= get_cosine_schedule_with_warmup(lr=0.0001, num_warmup_steps=4, num_training_steps=10)\n",
    "\n",
    "encoder_with_projection_head.summary()\n",
    "start = time.time()\n",
    "history = encoder_with_projection_head.fit(train_ds_c,validation_data=val_ds_c, batch_size=batch_size, \n",
    "                                           callbacks=[lr_schedule],#steps_per_epoch=19182 / batch_size,\n",
    "                                           epochs=10)\n",
    "\n",
    "print('duration run',(time.time()-start)/60/60,' hours')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rotate_tensor(image):\n",
    "    degree = tf.random.normal([])*360\n",
    "    image = tfa.image.rotate(image, degree * math.pi / 180)\n",
    "    return image\n",
    "\n",
    "def input_preprocess_train(image, label):\n",
    "    label = tf.one_hot(label, num_classes)\n",
    "    image = rotate_tensor(image)\n",
    "    image = tf.cast(image,'float32') #activate for densenet\n",
    "    image = preprocess_input(image)\n",
    "    return image, label\n",
    "\n",
    "def input_preprocess_val(image, label):\n",
    "    label = tf.one_hot(label, num_classes)\n",
    "    image = tf.cast(image,'float32') #activate for densenet\n",
    "    image = preprocess_input(image)\n",
    "    return image, label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds_s = train_ds.map(input_preprocess_train)\n",
    "val_ds_s = val_ds.map(input_preprocess_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_classifier3(encoder,  dropout, fc_layers):\n",
    "    inputs = keras.Input(shape=input_shape)\n",
    "    #for layer in encoder.layers:\n",
    "        #layer.trainable = True\n",
    "    encoder.trainable = True \n",
    "    for layer in encoder.layers[:-1]:\n",
    "        layer.trainable = False \n",
    "    x = encoder(inputs)\n",
    "    #x = layers.Flatten()(x)\n",
    "    for fc in fc_layers:\n",
    "        x = layers.Dense(fc, activation='relu')(x) # New FC layer, random init\n",
    "        x = layers.Dropout(dropout)(x)\n",
    "        \n",
    "    af_layer = ArcFace(output_dim=28, class_num=28, margin=0.5, scale=64.)  \n",
    "    arcface_output = af_layer(x)\n",
    "    model = tf.keras.Model(inputs,  arcface_output, name=\"grain-classifier\")\n",
    "    model.compile(optimizer='adam', loss=af_layer.loss, metrics=[acc_grain],run_eagerly=True)\n",
    "    return model\n",
    "\n",
    "def create_classifier(encoder):\n",
    "    #encoder.trainable = True\n",
    "    for layer in encoder.layers[:-1]:\n",
    "        layer.trainable = False\n",
    "\n",
    "    inputs = keras.Input(shape=input_shape)\n",
    "    features = encoder(inputs)\n",
    "    #features = layers.Dropout(dropout_rate)(features)\n",
    "    #features = layers.Dense(hidden_units, activation=\"relu\")(features)\n",
    "    #features = layers.Dropout(dropout_rate)(features)\n",
    "    outputs = layers.Dense(num_classes, activation=\"softmax\")(features)\n",
    "\n",
    "    model = keras.Model(inputs=inputs, outputs=outputs, name=\"cifar10-classifier\")\n",
    "    model.compile(\n",
    "        #optimizer=keras.optimizers.Adam(learning_rate),\n",
    "        optimizer='adam',\n",
    "        loss=\"categorical_crossentropy\",\n",
    "        metrics=[acc_grain],run_eagerly=True\n",
    "    )\n",
    "    return model\n",
    "\n",
    "def create_classifier2(encoder,  dropout, fc_layers):\n",
    "    inputs = keras.Input(shape=input_shape)\n",
    "    #for layer in encoder.layers:\n",
    "        #layer.trainable = True\n",
    "    encoder.trainable = True \n",
    "    for layer in encoder.layers[:-1]:\n",
    "        layer.trainable = False \n",
    "    x = encoder(inputs)\n",
    "    #x = layers.Flatten()(x)\n",
    "    for fc in fc_layers:\n",
    "        x = layers.Dense(fc, activation='relu')(x) # New FC layer, random init\n",
    "        x = layers.Dropout(dropout)(x)\n",
    "        \n",
    "    outputs = layers.Dense(num_classes, activation='softmax', name=\"pred\")(x) # New softmax layer\n",
    "    model = tf.keras.Model(inputs, outputs, name=\"cifar10-classifier\")\n",
    "    model.compile(optimizer='adam', loss=\"categorical_crossentropy\", metrics=[acc_grain],run_eagerly=True)\n",
    "    return model\n",
    "\n",
    "def build_model_attention(model):\n",
    "    inputs = keras.Input(shape=input_shape)\n",
    "    model.trainable = True\n",
    "    for layer in model.layers[:]:\n",
    "        layer.trainable = False\n",
    "    #model.trainable = False\n",
    "    pt_features = model(inputs)\n",
    "    pt_depth = model.get_output_shape_at(0)[-1]\n",
    "    bn_features = layers.BatchNormalization()(pt_features)\n",
    "    \n",
    "     # here we do an attention mechanism to turn pixels in the GAP on an off\n",
    "    attn_layer = layers.Conv2D(64, kernel_size = (1, 1), padding = 'same', activation = 'swish')(layers.Dropout(0.5)(bn_features))\n",
    "    attn_layer = layers.Conv2D(16, kernel_size = (1, 1), padding = 'same', activation = 'swish')(attn_layer)\n",
    "    attn_layer = layers.Conv2D(8, kernel_size = (1,1), padding = 'same', activation = 'swish')(attn_layer)\n",
    "    attn_layer = layers.Conv2D(1, kernel_size = (1, 1), padding = 'valid', activation = 'sigmoid')(attn_layer)\n",
    "    \n",
    "    # fan it out to all of the channels\n",
    "    up_c2_w = np.ones((1, 1, 1, pt_depth))\n",
    "    up_c2 = layers.Conv2D(pt_depth,kernel_size=(1,1),padding = 'same', activation='linear',use_bias = False,weights = [up_c2_w] )\n",
    "    up_c2.trainable = False\n",
    "    attn_layer = up_c2(attn_layer)\n",
    "    mask_features = layers.multiply([attn_layer, bn_features])\n",
    "    gap_features = layers.GlobalAveragePooling2D()(mask_features)\n",
    "    gap_mask = layers.GlobalAveragePooling2D()(attn_layer)\n",
    "    \n",
    "     # To account for missing values from the attention model\n",
    "    gap = layers.Lambda(lambda x: x[0] / x[1], name = 'RescaleGAP')([gap_features, gap_mask])\n",
    "    gap_dr = layers.Dropout(0.5)(gap)\n",
    "    dr_steps = layers.Dropout(0.25)(tf.keras.layers.Dense(128, activation = 'swish')(gap_dr))\n",
    "    \n",
    "    concat = layers.BatchNormalization()(dr_steps)\n",
    "    concat = layers.Dense(512, activation = 'swish')(concat)        \n",
    "    concat = layers.Dropout(0.15)(concat)\n",
    "    outputs = layers.Dense(28, activation=\"softmax\", name=\"pred\",dtype='float32')(concat)\n",
    "\n",
    "    model = tf.keras.Model([inputs], [outputs], name=\"EfficientNet\")\n",
    "    #optimizer = tf.keras.optimizers.Adam(learning_rate=1e-5)\n",
    "    model.compile(optimizer=\"adam\", loss=\"categorical_crossentropy\", metrics=[acc_grain],run_eagerly=True)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "encoder_with_projection_head.save('encoder_with_projection_head_resnet_bilinear.h5')\n",
    "#encoder_with_projection_head = keras.models.load_model('encoder_with_projection_head_mobilenet_bilinear.h5',custom_objects={'SupervisedContrastiveLoss':SupervisedContrastiveLoss},compile=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "#encoder_with_projection_head.save('encoder_with_projection_head_dense121')\n",
    "encoder_with_projection_head = keras.models.load_model('encoder_with_projection_head_dense121',custom_objects={'acc_grain':acc_grain},compile=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "model1 = tf.keras.Model(encoder_with_projection_head.get_layer('grains-encoder').input,encoder_with_projection_head.get_layer('grains-encoder').output)\n",
    "#model1= tf.keras.Model(model1.input, model1.get_layer('conv5_block3_3_conv').output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 00001: LearningRateScheduler reducing learning rate to 0.0.\n",
      "Epoch 1/10\n",
      "300/300 [==============================] - 264s 879ms/step - loss: 21.2570 - acc_grain: 0.0366 - val_loss: 3.3208 - val_acc_grain: 0.0207\n",
      "\n",
      "Epoch 00002: LearningRateScheduler reducing learning rate to 0.00025.\n",
      "Epoch 2/10\n",
      "300/300 [==============================] - 262s 875ms/step - loss: 15.2742 - acc_grain: 0.9418 - val_loss: 2.5314 - val_acc_grain: 0.9502\n",
      "\n",
      "Epoch 00003: LearningRateScheduler reducing learning rate to 0.0005.\n",
      "Epoch 3/10\n",
      "300/300 [==============================] - 262s 873ms/step - loss: 13.1406 - acc_grain: 0.9565 - val_loss: 2.3234 - val_acc_grain: 0.9663\n",
      "\n",
      "Epoch 00004: LearningRateScheduler reducing learning rate to 0.00075.\n",
      "Epoch 4/10\n",
      "300/300 [==============================] - 264s 880ms/step - loss: 12.3275 - acc_grain: 0.9616 - val_loss: 2.3085 - val_acc_grain: 0.9629\n",
      "\n",
      "Epoch 00005: LearningRateScheduler reducing learning rate to 0.001.\n",
      "Epoch 5/10\n",
      "300/300 [==============================] - 263s 877ms/step - loss: 12.2203 - acc_grain: 0.9600 - val_loss: 2.3086 - val_acc_grain: 0.9655\n",
      "\n",
      "Epoch 00006: LearningRateScheduler reducing learning rate to 0.0009330127018922195.\n",
      "Epoch 6/10\n",
      "300/300 [==============================] - 262s 873ms/step - loss: 12.2123 - acc_grain: 0.9613 - val_loss: 2.3100 - val_acc_grain: 0.9604\n",
      "\n",
      "Epoch 00007: LearningRateScheduler reducing learning rate to 0.00075.\n",
      "Epoch 7/10\n",
      "300/300 [==============================] - 262s 873ms/step - loss: 12.2098 - acc_grain: 0.9602 - val_loss: 2.3052 - val_acc_grain: 0.9622\n",
      "\n",
      "Epoch 00008: LearningRateScheduler reducing learning rate to 0.0005.\n",
      "Epoch 8/10\n",
      " 72/300 [======>.......................] - ETA: 2:42 - loss: 12.1706 - acc_grain: 0.9608"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-34-5711b3892b83>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m     13\u001b[0m \u001b[0mstart\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtime\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mtime\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     14\u001b[0m history = classifier.fit(train_ds_s, validation_data=val_ds_s, batch_size=batch_size, epochs=10,shuffle=True,\n\u001b[1;32m---> 15\u001b[1;33m                          \u001b[0mcallbacks\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mlr_schedule\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mverbose\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mclass_weight\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mweights\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     16\u001b[0m                         )\n\u001b[0;32m     17\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python36\\site-packages\\tensorflow\\python\\keras\\engine\\training.py\u001b[0m in \u001b[0;36m_method_wrapper\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m    106\u001b[0m   \u001b[1;32mdef\u001b[0m \u001b[0m_method_wrapper\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    107\u001b[0m     \u001b[1;32mif\u001b[0m \u001b[1;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_in_multi_worker_mode\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m  \u001b[1;31m# pylint: disable=protected-access\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 108\u001b[1;33m       \u001b[1;32mreturn\u001b[0m \u001b[0mmethod\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    109\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    110\u001b[0m     \u001b[1;31m# Running inside `run_distribute_coordinator` already.\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python36\\site-packages\\tensorflow\\python\\keras\\engine\\training.py\u001b[0m in \u001b[0;36mfit\u001b[1;34m(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_batch_size, validation_freq, max_queue_size, workers, use_multiprocessing)\u001b[0m\n\u001b[0;32m   1096\u001b[0m                 batch_size=batch_size):\n\u001b[0;32m   1097\u001b[0m               \u001b[0mcallbacks\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mon_train_batch_begin\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mstep\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 1098\u001b[1;33m               \u001b[0mtmp_logs\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtrain_function\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0miterator\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   1099\u001b[0m               \u001b[1;32mif\u001b[0m \u001b[0mdata_handler\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mshould_sync\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1100\u001b[0m                 \u001b[0mcontext\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0masync_wait\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python36\\site-packages\\tensorflow\\python\\keras\\engine\\training.py\u001b[0m in \u001b[0;36mtrain_function\u001b[1;34m(iterator)\u001b[0m\n\u001b[0;32m    804\u001b[0m       \u001b[1;32mdef\u001b[0m \u001b[0mtrain_function\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0miterator\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    805\u001b[0m         \u001b[1;34m\"\"\"Runs a training execution with one step.\"\"\"\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 806\u001b[1;33m         \u001b[1;32mreturn\u001b[0m \u001b[0mstep_function\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0miterator\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    807\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    808\u001b[0m     \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python36\\site-packages\\tensorflow\\python\\keras\\engine\\training.py\u001b[0m in \u001b[0;36mstep_function\u001b[1;34m(model, iterator)\u001b[0m\n\u001b[0;32m    793\u001b[0m         \u001b[1;32mreturn\u001b[0m \u001b[0moutputs\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    794\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 795\u001b[1;33m       \u001b[0mdata\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mnext\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0miterator\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    796\u001b[0m       \u001b[0moutputs\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mdistribute_strategy\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrun\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mrun_step\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0margs\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    797\u001b[0m       outputs = reduce_per_replica(\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python36\\site-packages\\tensorflow\\python\\data\\ops\\iterator_ops.py\u001b[0m in \u001b[0;36m__next__\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m    734\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    735\u001b[0m   \u001b[1;32mdef\u001b[0m \u001b[0m__next__\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m  \u001b[1;31m# For Python 3 compatibility\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 736\u001b[1;33m     \u001b[1;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mnext\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    737\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    738\u001b[0m   \u001b[1;32mdef\u001b[0m \u001b[0m_next_internal\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python36\\site-packages\\tensorflow\\python\\data\\ops\\iterator_ops.py\u001b[0m in \u001b[0;36mnext\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m    770\u001b[0m   \u001b[1;32mdef\u001b[0m \u001b[0mnext\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    771\u001b[0m     \u001b[1;32mtry\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 772\u001b[1;33m       \u001b[1;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_next_internal\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    773\u001b[0m     \u001b[1;32mexcept\u001b[0m \u001b[0merrors\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mOutOfRangeError\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    774\u001b[0m       \u001b[1;32mraise\u001b[0m \u001b[0mStopIteration\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python36\\site-packages\\tensorflow\\python\\data\\ops\\iterator_ops.py\u001b[0m in \u001b[0;36m_next_internal\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m    756\u001b[0m             \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_iterator_resource\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    757\u001b[0m             \u001b[0moutput_types\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_flat_output_types\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 758\u001b[1;33m             output_shapes=self._flat_output_shapes)\n\u001b[0m\u001b[0;32m    759\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    760\u001b[0m       \u001b[1;32mtry\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python36\\site-packages\\tensorflow\\python\\ops\\gen_dataset_ops.py\u001b[0m in \u001b[0;36miterator_get_next\u001b[1;34m(iterator, output_types, output_shapes, name)\u001b[0m\n\u001b[0;32m   2605\u001b[0m         \u001b[0m_ctx\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_context_handle\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtld\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mdevice_name\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m\"IteratorGetNext\"\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mname\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   2606\u001b[0m         \u001b[0mtld\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mop_callbacks\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0miterator\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m\"output_types\"\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0moutput_types\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 2607\u001b[1;33m         \"output_shapes\", output_shapes)\n\u001b[0m\u001b[0;32m   2608\u001b[0m       \u001b[1;32mreturn\u001b[0m \u001b[0m_result\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   2609\u001b[0m     \u001b[1;32mexcept\u001b[0m \u001b[0m_core\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_NotOkStatusException\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "#classifier = create_classifier(model1)\n",
    "#classifier = create_classifier2(model1, dropout=0.5, fc_layers=[512])\n",
    "classifier = create_classifier3(model1, dropout=0.5, fc_layers=[512])\n",
    "#classifier = build_model_attention(model1)\n",
    "#classifier = build_model_bilinear(model1)\n",
    "\n",
    "filename = 'best.h5'\n",
    "checkpoint = ModelCheckpoint(filename, monitor='val_acc_grain', verbose=1, save_best_only=True, mode='max')\n",
    "lr_schedule= get_cosine_schedule_with_warmup(lr=0.001, num_warmup_steps=4, num_training_steps=10)\n",
    "\n",
    "weights = dict(zip([0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27], [1,1,1,1,1,1,1,8.94,1,1,1,1,1,8.94,1,1,1,1,1,1,1,1,1,1,8.94,1,1,1]))\n",
    "\n",
    "start = time.time()\n",
    "history = classifier.fit(train_ds_s, validation_data=val_ds_s, batch_size=batch_size, epochs=10,shuffle=True,\n",
    "                         callbacks=[lr_schedule], verbose=1, class_weight=weights\n",
    "                        )\n",
    "\n",
    "print('duration run',(time.time()-start)/60/60,' hours')\n",
    "\n",
    "accuracy = classifier.evaluate(val_ds_s)[1]\n",
    "print(f\"Test accuracy: {round(accuracy * 100, 2)}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "ename": "SyntaxError",
     "evalue": "invalid syntax (<ipython-input-19-265111cba378>, line 1)",
     "output_type": "error",
     "traceback": [
      "\u001b[1;36m  File \u001b[1;32m\"<ipython-input-19-265111cba378>\"\u001b[1;36m, line \u001b[1;32m1\u001b[0m\n\u001b[1;33m    96.39%\u001b[0m\n\u001b[1;37m          ^\u001b[0m\n\u001b[1;31mSyntaxError\u001b[0m\u001b[1;31m:\u001b[0m invalid syntax\n"
     ]
    }
   ],
   "source": [
    "96.39%"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "# Plot training & validation iou_score values\n",
    "plt.figure(figsize=(30, 20))\n",
    "plt.subplot(121)\n",
    "plt.plot(history.history['acc_grain'])\n",
    "plt.plot(history.history['val_acc_grain'])\n",
    "plt.title('Model acc_score')\n",
    "plt.ylabel('acc_score')\n",
    "plt.xlabel('Epoch')\n",
    "plt.legend(['Train', 'Test'], loc='upper left')\n",
    "\n",
    "# Plot training & validation loss values\n",
    "plt.subplot(122)\n",
    "plt.plot(history.history['loss'])\n",
    "plt.plot(history.history['val_loss'])\n",
    "plt.title('Model loss')\n",
    "plt.ylabel('Loss')\n",
    "plt.xlabel('Epoch')\n",
    "plt.legend(['Train', 'Test'], loc='upper left')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 350
    },
    "executionInfo": {
     "elapsed": 2298,
     "status": "error",
     "timestamp": 1624373965288,
     "user": {
      "displayName": "ramdhan wibawa",
      "photoUrl": "",
      "userId": "14519448541243845830"
     },
     "user_tz": -420
    },
    "id": "UejZgfeXZ5C7",
    "outputId": "fbfcc699-a9ac-4b16-ccbb-0a8c4a25222e"
   },
   "outputs": [],
   "source": [
    "from tensorflow.keras.preprocessing import image\n",
    "import pandas as pd\n",
    "import zipfile\n",
    "import csv\n",
    "\n",
    "start = time.time()\n",
    "\n",
    "def load_class_list(class_list_file):\n",
    "    class_list = []\n",
    "    with open(class_list_file, 'r') as csvfile:\n",
    "        file_reader = csv.reader(csvfile)\n",
    "        for row in file_reader:\n",
    "            class_list.append(row[0])\n",
    "    class_list.sort()\n",
    "    return class_list\n",
    "\n",
    "class_list = load_class_list('class_list.txt')\n",
    "\n",
    "IMG_SIZE=224\n",
    "LBL = dict(zip(class_list, range(28)))\n",
    "\n",
    "cls_map = dict(zip(LBL.values(),LBL.keys()))\n",
    "\n",
    "def init_dir(pth):\n",
    "    if os.path.exists(pth):\n",
    "        shutil.rmtree(pth)\n",
    "    os.mkdir(pth)\n",
    "\n",
    "logger.info(\"loading the model\")\n",
    "model = model_fn(args['model_dir'])\n",
    "logger.info(\"Reading the test set\")\n",
    "\n",
    "res = []\n",
    "for fn in glob.iglob(args['data_dir'] + 'val/**/*.png', recursive=True):\n",
    "    file_name = os.path.basename(fn)\n",
    "    path = os.path.abspath(fn)\n",
    "    folder = os.path.split(os.path.dirname(path))[1]\n",
    "    if len(file_name.split(\"-\")) > 2:  # ignore master image with may grains, raw image names are in guid format\n",
    "        im = image.load_img(path, target_size=(IMG_SIZE, IMG_SIZE))\n",
    "        img_array = image.img_to_array(im)\n",
    "        img_batch = np.expand_dims(img_array, axis=0)\n",
    "        img_preprocessed = preprocess_input(img_batch)\n",
    "        pred = model.predict(img_preprocessed)\n",
    "        top3 = (-pred[0]).argsort()[:3]\n",
    "        res.append({'file_name': file_name, 'path': path, 'cls': folder, 'prediction':top3[0],  'proba_1':pred[0][top3[0]], 'prediction2':top3[1], 'proba_2':pred[0][top3[1]],  'prediction3':top3[2], 'proba_3':pred[0][top3[2]]})\n",
    "\n",
    "df = pd.DataFrame(res)\n",
    "df['prediction'] = df.prediction.map(cls_map)\n",
    "df['prediction2'] = df.prediction2.map(cls_map)\n",
    "df['prediction3'] = df.prediction3.map(cls_map)\n",
    "\n",
    "from sklearn.metrics import accuracy_score\n",
    "\n",
    "def w(x):\n",
    "    if \"_SOUND\" in x:\n",
    "        return 8.94 # weighting  Healthy grains more important than defective grains\n",
    "    return 1\n",
    "\n",
    "def scoring_fn(df):\n",
    "    \"\"\"\n",
    "    Weighted Accuracy Metric 90% Healthy grains, 10% unhealthy grains\n",
    "    \"\"\"\n",
    "    df['weight'] = df.cls.apply(w)\n",
    "    return accuracy_score(df.cls,df.prediction,sample_weight=df.weight)\n",
    "\n",
    "logger.info(f\"predictions have shape of {df.shape}\")\n",
    "\n",
    "score = scoring_fn(df)\n",
    "print (\"Accuracy Score:\" + str(score))\n",
    "# write to the output location\n",
    "\n",
    "print('Duration Run',(time.time()-start)/3600,' hours')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import classification_report\n",
    "print(classification_report(df['cls'], df['prediction']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from sklearn.metrics import confusion_matrix\n",
    "ticks=np.linspace(0, 27,num=28)\n",
    "plt.figure(figsize=(15,15))\n",
    "plt.imshow(confusion_matrix(df['cls'], df['prediction'],normalize='true'), interpolation='none')\n",
    "plt.colorbar()\n",
    "plt.xticks(ticks,fontsize=6)\n",
    "plt.yticks(ticks,fontsize=6)\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "id": "WZXP_9Kzfh6Q"
   },
   "outputs": [],
   "source": [
    "from tensorflow.keras.applications.resnet import preprocess_input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "s0bgHcc9rskI"
   },
   "outputs": [],
   "source": [
    "NUM_CLASSES=28\n",
    "model = build_model(num_classes=NUM_CLASSES)\n",
    "model.fit?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3Fcx6Qvnfh6Q",
    "outputId": "fc7bceb9-5ffb-46bc-ffaf-7e1cf07f05c0"
   },
   "outputs": [],
   "source": [
    "image.random_rotation(90)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "saWJd9X3fh6R"
   },
   "outputs": [],
   "source": [
    " rotation_range=args.rotation,\n",
    "      shear_range=args.shear,\n",
    "      zoom_range=args.zoom,\n",
    "      horizontal_flip=args.h_flip,\n",
    "      vertical_flip=args.v_flip"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "grain_unearth.ipynb",
   "provenance": []
  },
  "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}