{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Credits:\n* @marcosnovaes  https://www.kaggle.com/marcosnovaes/hubmap-looking-at-tfrecords and https://www.kaggle.com/marcosnovaes/hubmap-unet-keras-model-fit-with-tpu\n* @mgornergoogle https://www.kaggle.com/mgornergoogle/getting-started-with-100-flowers-on-tpu\n* qubvel https://github.com/qubvel/segmentation_models  !! 25 available backbones for each of 4 architectures\n"},{"metadata":{},"cell_type":"markdown","source":"## Setups and Imports"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%capture\n!pip install wandb -q\n# using https://github.com/qubvel/segmentation_models\n! pip install segmentation_models -q","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import tensorflow as tf\n\nimport os\nos.environ['SM_FRAMEWORK'] = 'tf.keras'\nimport glob\nimport numpy as np\nimport pandas as pd\nfrom sklearn.model_selection import KFold\nimport matplotlib\n%matplotlib inline\n\nfrom kaggle_datasets import KaggleDatasets\nGCS_PATH = KaggleDatasets().get_gcs_path('hubmap-tfrecord-512')\n\nprint(\"Tensorflow version \" + tf.__version__)\nAUTO = tf.data.experimental.AUTOTUNE\n\n# import segmentation models\nimport segmentation_models as sm\n\n# import W&B for ML experiment tracking\nimport wandb\nfrom wandb.keras import WandbCallback\n!wandb login 69f60a7711ce6b8bbae91ac6d15e45d6b1f1430e","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"#### Setup TPU"},{"metadata":{"trusted":true},"cell_type":"code","source":"try: # detect TPUs\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver() # TPU detection\n    tf.config.experimental_connect_to_cluster(tpu)\n    tf.tpu.experimental.initialize_tpu_system(tpu)\n    strategy = tf.distribute.experimental.TPUStrategy(tpu)\nexcept ValueError: # no TPU found, detect GPUs\n    #strategy = tf.distribute.MirroredStrategy() # for GPU or multi-GPU machines\n    strategy = tf.distribute.get_strategy() # default strategy that works on CPU and single GPU\n    #strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy() # for clusters of multi-GPU machines\n\nprint(\"Number of accelerators: \", strategy.num_replicas_in_sync)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Hyperparameters"},{"metadata":{"trusted":true},"cell_type":"code","source":"BATCH_SIZE = 8 * strategy.num_replicas_in_sync\nEPOCHS = 60\nBACKBONE = 'efficientnetb7' \nNFOLDS = 4\nSEED = 0\nVERBOSE = 1","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dataset\n\n### GCS_PATHS\n\nBased on: https://www.kaggle.com/marcosnovaes/hubmap-looking-at-tfrecords"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nuber_tile_df = pd.read_csv('/kaggle/input/hubmap-looking-at-tfrecords/train_all_tiles.csv')\nuber_tile_df['gcs_path'] = uber_tile_df.replace(regex = '/kaggle/input/hubmap-tfrecord-512',value = GCS_PATH)['local_path']\nuber_tile_df = uber_tile_df.loc[uber_tile_df['mask_density']  > 0].copy()\nuber_tile_df.shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Stratified folds"},{"metadata":{"trusted":true},"cell_type":"code","source":"img_ids = uber_tile_df['img_id'].unique()\n\nkf = KFold(n_splits=NFOLDS, random_state=SEED, shuffle=True)\nfor train_index, val_index in kf.split(img_ids): # one fold only currently\n    train_ids = [img_ids[ft] for ft in train_index]\n    val_ids = [img_ids[ft] for ft in val_index]\n    TRAINING_FILENAMES = list(uber_tile_df.loc[uber_tile_df['img_id'].isin(train_ids),'gcs_path'].values)\n    VALIDATION_FILENAMES = list(uber_tile_df.loc[uber_tile_df['img_id'].isin(val_ids),'gcs_path'].values)\n\nNUM_TRAINING_IMAGES = len(TRAINING_FILENAMES)\nNUM_VALIDATION_IMAGES = len(VALIDATION_FILENAMES)\nSTEPS_PER_EPOCH = NUM_TRAINING_IMAGES // BATCH_SIZE\nprint(NUM_VALIDATION_IMAGES)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Datasets pipeline"},{"metadata":{"trusted":true},"cell_type":"code","source":"# read back a record to make sure it the decoding works\npreprocess_input = sm.get_preprocessing(BACKBONE)\n\nimage_feature_description = {\n    'img_index': tf.io.FixedLenFeature([], tf.int64),\n    'height': tf.io.FixedLenFeature([], tf.int64),\n    'width': tf.io.FixedLenFeature([], tf.int64),\n    'num_channels': tf.io.FixedLenFeature([], tf.int64),\n    'img_bytes': tf.io.FixedLenFeature([], tf.string),\n    'mask': tf.io.FixedLenFeature([], tf.string),\n    'tile_id': tf.io.FixedLenFeature([], tf.int64),\n    'tile_col_pos': tf.io.FixedLenFeature([], tf.int64),\n    'tile_row_pos': tf.io.FixedLenFeature([], tf.int64),\n}\n\ndef _parse_image_function(example_proto):\n    single_example = tf.io.parse_single_example(example_proto, image_feature_description)\n    image = tf.reshape( tf.io.decode_raw(single_example['img_bytes'],out_type='uint8'), (512, 512, 3))\n    image = tf.image.convert_image_dtype(image, tf.float32)\n    \n    mask =  tf.reshape(tf.io.decode_raw(single_example['mask'],out_type='bool'),(512, 512,1))\n\n    return image, tf.cast(mask, tf.float32) # cast as float32 required for TPU\n\ndef load_dataset(filenames, ordered=False):\n    ignore_order = tf.data.Options()\n    if not ordered:\n        ignore_order.experimental_deterministic = False\n    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTO, compression_type=\"GZIP\")\n    dataset = dataset.with_options(ignore_order)\n    dataset = dataset.map(_parse_image_function, num_parallel_calls=AUTO)\n    return dataset\n\ndef get_training_dataset():\n    dataset = load_dataset(TRAINING_FILENAMES)\n    dataset = dataset.repeat()\n    dataset = dataset.shuffle(128)\n    dataset = dataset.batch(BATCH_SIZE,drop_remainder=True)\n    dataset = dataset.prefetch(AUTO)\n    return dataset\n\ndef get_validation_dataset(ordered=False):\n    dataset = load_dataset(VALIDATION_FILENAMES, ordered=ordered)\n    dataset = dataset.batch(BATCH_SIZE,drop_remainder=True)\n    dataset = dataset.cache()\n    dataset = dataset.prefetch(AUTO)\n    return dataset","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"with strategy.scope():\n    model = sm.Unet(BACKBONE)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Callbacks"},{"metadata":{},"cell_type":"markdown","source":"### Early Stopping"},{"metadata":{"trusted":true},"cell_type":"code","source":"early_stopper = tf.keras.callbacks.EarlyStopping(monitor='val_loss', \n                                                 patience=10, mode='min')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Model Checkpoint"},{"metadata":{"trusted":true},"cell_type":"code","source":"model_ckpt = tf.keras.callbacks.ModelCheckpoint(filepath='model_weights.h5', \n                                                monitor='val_loss', \n                                                save_weights_only=True,\n                                                save_best_only=True, \n                                                mode='min')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Custom Callback to Visualize Segmentation Masks using W&B"},{"metadata":{"trusted":true},"cell_type":"code","source":"segmentation_classes = ['issue', 'no issue']\n\n# returns a dictionary of labels\ndef labels():\n  l = {}\n  for i, label in enumerate(segmentation_classes):\n    l[i] = label\n  return l\n\n# util function for generating interactive image mask from components\ndef wandb_mask(bg_img, pred_mask, true_mask):\n  return wandb.Image(bg_img, masks={\n      \"prediction\" : {\n          \"mask_data\" : pred_mask, \n          \"class_labels\" : labels()\n      },\n      \"ground truth\" : {\n          \"mask_data\" : true_mask, \n          \"class_labels\" : labels()\n      }\n    }\n  )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class SemanticLogger(tf.keras.callbacks.Callback):\n    def __init__(self, dataloader):\n        super(SemanticLogger, self).__init__()\n        self.val_images, self.val_masks = next(iter(dataloader))\n\n    def on_epoch_end(self, logs, epoch):\n        pred_masks = self.model.predict(self.val_images)\n        pred_masks = np.argmax(pred_masks, axis=-1)\n\n        val_images = tf.image.convert_image_dtype(self.val_images, tf.uint8)\n        val_masks = tf.image.convert_image_dtype(self.val_masks, tf.uint8)\n        val_masks = tf.squeeze(val_masks, axis=-1)\n        \n        pred_masks = tf.image.convert_image_dtype(pred_masks, tf.uint8)\n\n        mask_list = []\n        for i in range(len(self.val_images)):\n          mask_list.append(wandb_mask(val_images[i].numpy(), \n                                      pred_masks[i].numpy(), \n                                      val_masks[i].numpy()))\n\n        wandb.log({\"predictions\" : mask_list})","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train Model Using W&B"},{"metadata":{"trusted":true},"cell_type":"code","source":"# compile model\noptimizer = 'adam'\nmodel.compile(optimizer=optimizer,\n              loss=tf.keras.losses.BinaryCrossentropy(),    \n              metrics=[sm.metrics.iou_score,'accuracy'])\n\n# initialize wandb run\nwandb.init(project='HuBMAP')\n\n_ = model.fit(get_training_dataset(), \n              epochs=EPOCHS,\n              steps_per_epoch=STEPS_PER_EPOCH,\n              verbose = VERBOSE,\n              validation_data=get_validation_dataset(),\n              callbacks=[early_stopper,\n                         model_ckpt,\n                         WandbCallback(),\n                         SemanticLogger(get_validation_dataset())\n                        ])\n\nwandb.finish()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# save whole model for submission without internet\nmodel.load_weights('model_weights.h5')\nmodel.save('model.h5')","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}