{"cells":[{"metadata":{},"cell_type":"markdown","source":"# [Training Kernel](https://www.kaggle.com/durbin164/tpu-visual-transformer-vit-keras-tf)"},{"metadata":{},"cell_type":"markdown","source":"## Install ViT-keras"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"!pip install ../input/validators\n!cp -R ../input/vit-keras/vit_keras ./","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Imports requrements"},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nimport matplotlib.pyplot as plt\nimport pandas as pd\nimport seaborn as sns\nimport os\nimport cv2\nimport re\nimport math\n\nfrom sklearn.model_selection import KFold\nfrom sklearn.metrics import confusion_matrix, classification_report\n\nfrom tensorflow.keras.preprocessing.image import ImageDataGenerator\nfrom sklearn.model_selection import train_test_split\nfrom tensorflow.keras.losses import CategoricalCrossentropy\nimport tensorflow as tf\n\nfrom tensorflow.keras.layers import Dense, Dropout,\\\n        Flatten,GlobalAveragePooling2D,BatchNormalization, Activation\n\nfrom tensorflow.keras.models import Sequential, Model\nfrom tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping, ReduceLROnPlateau\nfrom tensorflow.keras import layers\n\nimport tensorflow.keras.layers as L\nimport tensorflow.keras.backend as K\n\nfrom kaggle_datasets import KaggleDatasets\n\n# import efficientnet.keras as efn \nfrom vit_keras import vit, utils\nprint(tf.__version__)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"# Detect hardware, return appropriate distribution strategy\ntry:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n    print(f'Running on TPU {tpu.master()}')\nexcept ValueError:\n    tpu = None\n\nif tpu:\n    tf.config.experimental_connect_to_cluster(tpu)\n    tf.tpu.experimental.initialize_tpu_system(tpu)\n    strategy = tf.distribute.experimental.TPUStrategy(tpu)\nelse:\n    strategy = tf.distribute.get_strategy()\n\nAUTO = tf.data.experimental.AUTOTUNE\nREPLICAS = strategy.num_replicas_in_sync\nprint(f'REPLICAS: {REPLICAS}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#setting\nSEED = 100\nDEBUG = False\nWANDB = False\n# TARGET_SIZE = 512\nVALIDATION_SIZE = 0.2\nBATCH_SIZE = 8 *REPLICAS\nLEARNING_RATE = 3e-5 * REPLICAS\nEPOCHS=40\nMODEL_NAME = \"VitL16\"\nN_FOLDS = 5\n\nTTA = False\nN_TTA = 3\n\n#For BiTempered loss function\nT_1 = 0.2\nT_2 = 1.2\nSMOOTH_FRACTION = 0.01\nN_ITER = 5\n\nHEIGHT = 512\nWIDTH = 512\nHEIGHT_RS = 512\nWIDTH_RS = 512\nCHANNELS = 3\nN_CLASSES = 5\n\nMODEL_SAVE_PATH = \"\"\n\nif DEBUG:\n    EPOCHS = 2\n    # LEARNING_RATE = 3e-5 * REPLICAS\n    # TARGET_SIZE = 512\n    # BATCH_SIZE = 8*REPLICAS","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Data augmentations"},{"metadata":{"trusted":true},"cell_type":"code","source":"# data augmentation @cdeotte kernel: https://www.kaggle.com/cdeotte/rotation-augmentation-gpu-tpu-0-96\ndef transform_rotation(image, height, rotation):\n    # input image - is one image of size [dim,dim,3] not a batch of [b,dim,dim,3]\n    # output - image randomly rotated\n    DIM = height\n    XDIM = DIM%2 #fix for size 331\n    \n    rotation = rotation * tf.random.uniform([1],dtype='float32')\n    # CONVERT DEGREES TO RADIANS\n    rotation = math.pi * rotation / 180.\n    \n    # ROTATION MATRIX\n    c1 = tf.math.cos(rotation)\n    s1 = tf.math.sin(rotation)\n    one = tf.constant([1],dtype='float32')\n    zero = tf.constant([0],dtype='float32')\n    rotation_matrix = tf.reshape(tf.concat([c1,s1,zero, -s1,c1,zero, zero,zero,one],axis=0),[3,3])\n\n    # LIST DESTINATION PIXEL INDICES\n    x = tf.repeat( tf.range(DIM//2,-DIM//2,-1), DIM )\n    y = tf.tile( tf.range(-DIM//2,DIM//2),[DIM] )\n    z = tf.ones([DIM*DIM],dtype='int32')\n    idx = tf.stack( [x,y,z] )\n    \n    # ROTATE DESTINATION PIXELS ONTO ORIGIN PIXELS\n    idx2 = K.dot(rotation_matrix,tf.cast(idx,dtype='float32'))\n    idx2 = K.cast(idx2,dtype='int32')\n    idx2 = K.clip(idx2,-DIM//2+XDIM+1,DIM//2)\n    \n    # FIND ORIGIN PIXEL VALUES \n    idx3 = tf.stack( [DIM//2-idx2[0,], DIM//2-1+idx2[1,]] )\n    d = tf.gather_nd(image, tf.transpose(idx3))\n        \n    return tf.reshape(d,[DIM,DIM,3])\n\ndef transform_shear(image, height, shear):\n    # input image - is one image of size [dim,dim,3] not a batch of [b,dim,dim,3]\n    # output - image randomly sheared\n    DIM = height\n    XDIM = DIM%2 #fix for size 331\n    \n    shear = shear * tf.random.uniform([1],dtype='float32')\n    shear = math.pi * shear / 180.\n        \n    # SHEAR MATRIX\n    one = tf.constant([1],dtype='float32')\n    zero = tf.constant([0],dtype='float32')\n    c2 = tf.math.cos(shear)\n    s2 = tf.math.sin(shear)\n    shear_matrix = tf.reshape(tf.concat([one,s2,zero, zero,c2,zero, zero,zero,one],axis=0),[3,3])    \n\n    # LIST DESTINATION PIXEL INDICES\n    x = tf.repeat( tf.range(DIM//2,-DIM//2,-1), DIM )\n    y = tf.tile( tf.range(-DIM//2,DIM//2),[DIM] )\n    z = tf.ones([DIM*DIM],dtype='int32')\n    idx = tf.stack( [x,y,z] )\n    \n    # ROTATE DESTINATION PIXELS ONTO ORIGIN PIXELS\n    idx2 = K.dot(shear_matrix,tf.cast(idx,dtype='float32'))\n    idx2 = K.cast(idx2,dtype='int32')\n    idx2 = K.clip(idx2,-DIM//2+XDIM+1,DIM//2)\n    \n    # FIND ORIGIN PIXEL VALUES \n    idx3 = tf.stack( [DIM//2-idx2[0,], DIM//2-1+idx2[1,]] )\n    d = tf.gather_nd(image, tf.transpose(idx3))\n        \n    return tf.reshape(d,[DIM,DIM,3])\n\n# CutOut\ndef data_augment_cutout(image, min_mask_size=(int(HEIGHT * .1), int(HEIGHT * .1)), \n                        max_mask_size=(int(HEIGHT * .125), int(HEIGHT * .125))):\n    p_cutout = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    \n    if p_cutout > .85: # 10~15 cut outs\n        n_cutout = tf.random.uniform([], 10, 15, dtype=tf.int32)\n        image = random_cutout(image, HEIGHT, WIDTH, \n                              min_mask_size=min_mask_size, max_mask_size=max_mask_size, k=n_cutout)\n    elif p_cutout > .6: # 5~10 cut outs\n        n_cutout = tf.random.uniform([], 5, 10, dtype=tf.int32)\n        image = random_cutout(image, HEIGHT, WIDTH, \n                              min_mask_size=min_mask_size, max_mask_size=max_mask_size, k=n_cutout)\n    elif p_cutout > .25: # 2~5 cut outs\n        n_cutout = tf.random.uniform([], 2, 5, dtype=tf.int32)\n        image = random_cutout(image, HEIGHT, WIDTH, \n                              min_mask_size=min_mask_size, max_mask_size=max_mask_size, k=n_cutout)\n    else: # 1 cut out\n        image = random_cutout(image, HEIGHT, WIDTH, \n                              min_mask_size=min_mask_size, max_mask_size=max_mask_size, k=1)\n\n    return image\n\ndef random_cutout(image, height, width, channels=3, min_mask_size=(10, 10), max_mask_size=(80, 80), k=1):\n    assert height > min_mask_size[0]\n    assert width > min_mask_size[1]\n    assert height > max_mask_size[0]\n    assert width > max_mask_size[1]\n\n    for i in range(k):\n      mask_height = tf.random.uniform(shape=[], minval=min_mask_size[0], maxval=max_mask_size[0], dtype=tf.int32)\n      mask_width = tf.random.uniform(shape=[], minval=min_mask_size[1], maxval=max_mask_size[1], dtype=tf.int32)\n\n      pad_h = height - mask_height\n      pad_top = tf.random.uniform(shape=[], minval=0, maxval=pad_h, dtype=tf.int32)\n      pad_bottom = pad_h - pad_top\n\n      pad_w = width - mask_width\n      pad_left = tf.random.uniform(shape=[], minval=0, maxval=pad_w, dtype=tf.int32)\n      pad_right = pad_w - pad_left\n\n      cutout_area = tf.zeros(shape=[mask_height, mask_width, channels], dtype=tf.uint8)\n\n      cutout_mask = tf.pad([cutout_area], [[0,0],[pad_top, pad_bottom], [pad_left, pad_right], [0,0]], constant_values=1)\n      cutout_mask = tf.squeeze(cutout_mask, axis=0)\n      image = tf.multiply(tf.cast(image, tf.float32), tf.cast(cutout_mask, tf.float32))\n\n    return image","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def data_augment(image, label):\n    p_rotation = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_spatial = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_rotate = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_pixel_1 = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_pixel_2 = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_pixel_3 = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_shear = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_crop = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_cutout = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    flip_left = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    flip_right = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    \n#     # Shear\n#     if p_shear > .2:\n#         if p_shear > .6:\n#             image = transform_shear(image, HEIGHT, shear=20.)\n#         else:\n#             image = transform_shear(image, HEIGHT, shear=-20.)\n            \n    # Rotation\n    if p_rotation > .2:\n        if p_rotation > .6:\n            image = transform_rotation(image, HEIGHT, rotation=10.)\n        else:\n            image = transform_rotation(image, HEIGHT, rotation=-10.)\n            \n    # Flips\n    if flip_left>0.5:\n        image = tf.image.random_flip_left_right(image)\n    if flip_right>0.5 :\n        image = tf.image.random_flip_up_down(image)\n#     if p_spatial > .75:\n#         image = tf.image.transpose(image)\n        \n#     # Rotates\n#     if p_rotate > .75:\n#         image = tf.image.rot90(image, k=3) # rotate 270º\n#     elif p_rotate > .5:\n#         image = tf.image.rot90(image, k=2) # rotate 180º\n#     elif p_rotate > .25:\n#         image = tf.image.rot90(image, k=1) # rotate 90º\n        \n#     # Pixel-level transforms\n#     if p_pixel_1 >= .4:\n#         image = tf.image.random_saturation(image, lower=.7, upper=1.3)\n#     if p_pixel_2 >= .4:\n#         image = tf.image.random_contrast(image, lower=.8, upper=1.2)\n#     if p_pixel_3 >= .4:\n#         image = tf.image.random_brightness(image, max_delta=.1)\n        \n#     # Crops\n#     if p_crop > .6:\n#         if p_crop > .9:\n#             image = tf.image.central_crop(image, central_fraction=.5)\n#         elif p_crop > .8:\n#             image = tf.image.central_crop(image, central_fraction=.6)\n#         elif p_crop > .7:\n#             image = tf.image.central_crop(image, central_fraction=.7)\n#         else:\n#             image = tf.image.central_crop(image, central_fraction=.8)\n#     elif p_crop > .3:\n#         crop_size = tf.random.uniform([], int(HEIGHT*.6), HEIGHT, dtype=tf.int32)\n#         image = tf.image.random_crop(image, size=[crop_size, crop_size, CHANNELS])\n            \n    image = tf.image.resize(image, size=[HEIGHT, WIDTH])\n\n#     if p_cutout > .5:\n#         image = data_augment_cutout(image)\n        \n    return image, label","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Data loder"},{"metadata":{"trusted":true},"cell_type":"code","source":"def center_crop(image):\n#     image = tf.reshape(image, [600, 800, CHANNELS]) # Original shape\n    image = tf.image.resize(image, [600, 800])\n    \n    h, w = image.shape[0], image.shape[1]\n    if h > w:\n        image = tf.image.crop_to_bounding_box(image, (h - w) // 2, 0, w, w)\n    else:\n        image = tf.image.crop_to_bounding_box(image, 0, (w - h) // 2, h, h)\n        \n#     image = tf.image.resize(image, [HEIGHT, WIDTH]) # Expected shape\n    return image\n\n# Datasets utility functions\ndef decode_image(image_data):\n    \"\"\"\n        1. Decode a JPEG-encoded image to a uint8 tensor.\n        2. Cast tensor to float and normalizes (range between 0 and 1).\n        3. Resize and reshape images to the expected size.\n    \"\"\"\n    image = tf.image.decode_jpeg(image_data, channels=3)\n    \n    image = tf.cast(image, tf.float32)/255.0\n    \n#     image = center_crop(image)\n    \n    \n   \n                      \n    image = tf.image.resize(image, [HEIGHT, WIDTH])\n    image = tf.reshape(image, [HEIGHT, WIDTH, 3])\n    return image\n\ndef read_tfrecord(example, labeled=True):\n    \"\"\"\n        1. Parse data based on the 'TFREC_FORMAT' map.\n        2. Decode image.\n        3. If 'labeled' returns (image, label) if not (image, name).\n    \"\"\"\n    if labeled:\n        TFREC_FORMAT = {\n            'image': tf.io.FixedLenFeature([], tf.string), \n            'target': tf.io.FixedLenFeature([], tf.int64), \n        }\n    else:\n        TFREC_FORMAT = {\n            'image': tf.io.FixedLenFeature([], tf.string), \n            'image_name': tf.io.FixedLenFeature([], tf.string), \n        }\n    example = tf.io.parse_single_example(example, TFREC_FORMAT)\n    image = decode_image(example['image'])\n   \n    if labeled:\n        label_or_name = tf.cast(example['target'], tf.int32)\n    else:\n        label_or_name =  example['image_name']\n    return image, label_or_name\n\ndef load_dataset(filenames, labeled=True, ordered=False):\n    \"\"\"\n        Create a Tensorflow dataset from TFRecords.\n    \"\"\"\n    ignore_order = tf.data.Options()\n    if not ordered:\n        ignore_order.experimental_deterministic = False\n\n    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTO)\n    dataset = dataset.with_options(ignore_order)\n    dataset = dataset.map(lambda x: read_tfrecord(x, labeled=labeled), num_parallel_calls=AUTO)\n    return dataset\n\ndef get_dataset(FILENAMES, labeled=True, ordered=False, repeated=False, augment=False):\n    \"\"\"\n        Return a Tensorflow dataset ready for training or inference.\n    \"\"\"\n    dataset = load_dataset(FILENAMES, labeled=labeled, ordered=ordered)\n    if augment:\n        dataset = dataset.map(data_augment, num_parallel_calls=AUTO)\n    if repeated:\n        dataset = dataset.repeat()\n    if not ordered:\n        dataset = dataset.shuffle(2048)\n    dataset = dataset.batch(BATCH_SIZE)\n    dataset = dataset.prefetch(AUTO)\n    return dataset","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Bitemperate loss"},{"metadata":{"trusted":true},"cell_type":"code","source":"# import module we'll need to import our custom module\nfrom shutil import copyfile\n\n# copy our file into the working directory (make sure it has .py suffix)\ncopyfile(src = \"../input/bitempered-logistic-loss-tensorflow-v2/bi_tempered_loss.py\", dst = \"../working/loss.py\")\n\n# import all our functions\nfrom loss import bi_tempered_logistic_loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"with strategy.scope():\n  class BiTemperedLogisticLoss(tf.keras.losses.Loss):\n    def __init__(self, t1, t2, lbl_smth, n_iter):\n      super(BiTemperedLogisticLoss, self).__init__()\n      self.t1 = t1\n      self.t2 = t2\n      self.lbl_smth = lbl_smth\n      self.n_iter = n_iter\n\n    def call(self, y_true, y_pred):\n      return bi_tempered_logistic_loss(y_pred, y_true, self.t1, self.t2, self.lbl_smth, self.n_iter)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Create model"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_vit_model(weights = None):\n    if weights == None:\n       return   vit.vit_l16(\n                      image_size=HEIGHT,\n                      activation='softmax',\n                      pretrained=False,\n                      include_top=True,\n                      pretrained_top=False,\n                      classes = 5,\n#                       weights = 'imagenet21k'\n       )\n\n    else:\n        return vit.vit_l16(\n                      image_size=HEIGHT,\n                      activation='softmax',\n                      pretrained=False,\n                      include_top=True,\n                      pretrained_top=False,\n                      classes = 5,\n                      weights = weights\n       )\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Predict"},{"metadata":{"trusted":true},"cell_type":"code","source":"# GCS_PATH= KaggleDatasets().get_gcs_path('cassava-leaf-disease-classification')+'/test_tfrecords' \nfiles_path = '../input/cassava-leaf-disease-classification/test_images'\n\nTEST_FILENAMES = tf.io.gfile.glob('../input/cassava-leaf-disease-classification/test_tfrecords/*')\n# print(TEST_FILENAMES)\n# model_base_path = '../input/cassava-efficientnetb4'\n\n# model_path_list = os.listdir(model_base_path)\n# print(model_path_list)\n\n# model_path_list = [x for x in model_path_list if x.startswith('EfficentNet4_best_btl_') ]\n\nmodel_path_list = [\n#     '../input/tpu-vit-keras/ViTB16_best_fold_0_.h5',\n#     '../input/cassava-leaf-vit-models/ViTL16_best_fold_0_v_1.h5',\n#     '../input/cassava-leaf-vit-models/ViTB16_best_fold_0_v_2.h5',\n    '../input/cassava-leaf-vit-models/ViTB16_best_fold_0__v3_.h5',\n    '../input/cassava-leaf-vit-models/ViTB16_best_fold_0_v4_.h5',\n    '../input/cassava-leaf-vit-models/ViTB16_best_fold_1_v4_.h5',\n    '../input/cassava-leaf-vit-models/ViTB16_best_fold_2_v4_.h5',\n    '../input/cassava-leaf-vit-models/ViTB16_best_fold_3_v4_.h5'\n                   \n                  ]\n\ntest_preds = np.zeros((len(os.listdir(files_path)), N_CLASSES))\n\nfor idx, model_path in enumerate(model_path_list):\n    print(\"Model: \", model_path)\n    tf.keras.backend.clear_session()\n    \n#     if idx == 2:\n    model = get_vit_model('imagenet21k')\n    model.load_weights(model_path)\n#     else:\n#         model = get_vit_model()\n#         model.load_weights(model_path)\n    \n    if TTA:\n        \n        for step in range(N_TTA):\n            print(f\"TTA step: {step+1}/{N_TTA}\")\n            test_ds = get_dataset(TEST_FILENAMES, labeled=False, ordered=True, augment = True)\n            x_test = test_ds.map(lambda image, image_name: image)\n            test_preds += model.predict(x_test)\n        \n    \n    else:\n        test_ds = get_dataset(TEST_FILENAMES, labeled=False, ordered=True)\n        x_test = test_ds.map(lambda image, image_name: image)\n        test_preds += model.predict(x_test)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_preds = np.argmax(test_preds, axis=-1)\nimage_names = [img_name.numpy().decode('utf-8') for img, img_name in iter(test_ds.unbatch())]\n\nsubmission = pd.DataFrame({'image_id': image_names, 'label': test_preds})\nsubmission.to_csv('submission.csv', index=False)\ndisplay(submission.head())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}