train_pretrained_segmentation_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialize and load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "# Necessary for CUDA 10 or something?\n",
    "os.environ[\"TF_FORCE_GPU_ALLOW_GROWTH\"] = \"true\"\n",
    "#os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = \"1\"\n",
    "#os.environ[\"TF_ENABLE_AUTO_MIXED_PRECISION_GRAPH_REWRITE\"] = \"1\"\n",
    "#os.environ[\"TF_ENABLE_AUTO_MIXED_PRECISION_LOSS_SCALING\"] = \"1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml\n",
    "import random\n",
    "import json\n",
    "from datetime import datetime\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from steel_seg.utils import (\n",
    "    dice_coeff_kaggle,\n",
    "    per_class_dice_coeff,\n",
    "    rle_to_dense,\n",
    "    dense_to_rle,\n",
    "    visualize_segmentations,\n",
    "    onehottify)\n",
    "from steel_seg.dataset.severstal_steel_dataset_generator import \\\n",
    "    SeverstalSteelDatasetGenerator\n",
    "from steel_seg.dataset.dataset_utils import load_annotations, split_data\n",
    "from steel_seg.train import (\n",
    "    class_weighted_binary_crossentropy,\n",
    "    weighted_binary_crossentropy,\n",
    "    pixel_map_weighted_binary_crossentropy,\n",
    "    dice_loss_multi_class,\n",
    "    dice_coef,\n",
    "    DiceCoefByClassAndEmptiness,\n",
    "    eval)\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('SETTINGS.yaml') as f:\n",
    "    cfg = yaml.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "anns_dict = load_annotations(cfg['TRAIN_ANNOTATIONS_FILE'])\n",
    "imgs = list(anns_dict.keys())\n",
    "\n",
    "test_imgs, val_imgs, train_imgs = split_data(imgs,\n",
    "                                             test_split=cfg['TEST_SPLIT'],\n",
    "                                             val_split=cfg['VAL_SPLIT'],\n",
    "                                             batch_size=cfg['SEGMENTATION_BATCH_SIZE'],\n",
    "                                             load_cached=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = SeverstalSteelDatasetGenerator(\n",
    "     train_imgs,\n",
    "     is_training=True,\n",
    "     train_img_dir=cfg['TRAIN_IMAGE_DIR'],\n",
    "     train_anns_file=cfg['TRAIN_ANNOTATIONS_FILE'],\n",
    "     img_height=cfg['IMG_HEIGHT'],\n",
    "     img_width=cfg['IMG_WIDTH'],\n",
    "     num_classes=cfg['NUM_CLASSES'],\n",
    "     batch_size=cfg['SEGMENTATION_BATCH_SIZE'],\n",
    "     brightness_max_delta=cfg['BRIGHTNESS_MAX_DELTA'],\n",
    "     contrast_lower_factor=cfg['CONTRAST_LOWER_FACTOR'],\n",
    "     contrast_upper_factor=cfg['CONTRAST_UPPER_FACTOR'],\n",
    "     balance_classes=cfg['SEGMENTATION_BALANCE_CLASSES'],\n",
    "     max_oversample_rate=cfg['SEGMENTATION_MAX_OVERSAMPLE_RATE'])\n",
    "\n",
    "val_data = SeverstalSteelDatasetGenerator(\n",
    "     val_imgs,\n",
    "     is_training=False,\n",
    "     train_img_dir=cfg['TRAIN_IMAGE_DIR'],\n",
    "     train_anns_file=cfg['TRAIN_ANNOTATIONS_FILE'],\n",
    "     img_height=cfg['IMG_HEIGHT'],\n",
    "     img_width=cfg['IMG_WIDTH'],\n",
    "     num_classes=cfg['NUM_CLASSES'],\n",
    "     batch_size=cfg['SEGMENTATION_BATCH_SIZE'],\n",
    "     brightness_max_delta=None,\n",
    "     contrast_lower_factor=None,\n",
    "     contrast_upper_factor=None,\n",
    "     balance_classes=None,\n",
    "     max_oversample_rate=None)\n",
    "\n",
    "test_data = SeverstalSteelDatasetGenerator(\n",
    "     test_imgs,\n",
    "     is_training=False,\n",
    "     train_img_dir=cfg['TRAIN_IMAGE_DIR'],\n",
    "     train_anns_file=cfg['TRAIN_ANNOTATIONS_FILE'],\n",
    "     img_height=cfg['IMG_HEIGHT'],\n",
    "     img_width=cfg['IMG_WIDTH'],\n",
    "     num_classes=cfg['NUM_CLASSES'],\n",
    "     batch_size=1,\n",
    "     brightness_max_delta=None,\n",
    "     contrast_lower_factor=None,\n",
    "     contrast_upper_factor=None,\n",
    "     balance_classes=None,\n",
    "     max_oversample_rate=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# count = 0\n",
    "# for ex in train_data:\n",
    "#     count += 1\n",
    "#     if count == 4:\n",
    "#         break\n",
    "# img, ann = ex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# c = 1\n",
    "\n",
    "# plt.figure(figsize=(20, 10))\n",
    "# plt.imshow(np.repeat(img[c, :, :, :], 3, axis=-1))\n",
    "# plt.show()\n",
    "# for i in range(4):\n",
    "#     plt.figure(figsize=(20, 10))\n",
    "#     plt.imshow(ann[c, :, :, i])\n",
    "#     plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(test_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#os.environ['SM_FRAMEWORK'] = 'tf.keras'\n",
    "import keras\n",
    "import segmentation_models as sm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.keras.backend.set_image_data_format('channels_last')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "backbone = 'resnet34'#'mobilenetv2' # 'resnet34' #https://github.com/qubvel/segmentation_models#Backbones\n",
    "base_model = sm.Unet(backbone, classes=cfg['NUM_CLASSES'], encoder_weights='imagenet', activation='sigmoid')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO confirm that this is the correct preprocessing\n",
    "input = keras.Input(shape=(cfg['IMG_HEIGHT'], cfg['IMG_WIDTH'], 1))\n",
    "# Necessary to wrap in keras.layers.Lambda so that save_model works\n",
    "x = keras.layers.Lambda(lambda x: tf.tile(x / 127.5 - 1.0, [1, 1, 1, 3]))(input)\n",
    "output = base_model(x)\n",
    "seg_model = keras.Model(inputs=[input], outputs=[output])\n",
    "\n",
    "#model_checkpoint_name = 'resnet34_unet_pretrained_imgaug'\n",
    "model_checkpoint_name = 'resnet34_unet_imgaug_focal'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cls_weights = [1.0, 1.0, 1.0, 1.0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = cfg['SEGMENTATION_BATCH_SIZE']\n",
    "\n",
    "# set class weights for dice_loss (0, 1, 2, 3, background)\n",
    "#dice_loss = sm.losses.DiceLoss(class_weights=np.array([1.0, 1.0, 1.0, 1.0, 1.0]))\n",
    "# JaccardLoss\n",
    "# See https://github.com/qubvel/segmentation_models/blob/master/examples/multiclass%20segmentation%20(camvid).ipynb\n",
    "# See: https://github.com/qubvel/segmentation_models/blob/master/segmentation_models/losses.py\n",
    "\n",
    "\n",
    "seg_model.compile(\n",
    "    optimizer=tf.train.AdamOptimizer(0.00001),\n",
    "    loss=pixel_map_weighted_binary_crossentropy(cls_weights), #sm.losses.categorical_focal_loss, #dice_loss_multi_class,\n",
    "    metrics=[\n",
    "        tf.keras.metrics.BinaryAccuracy(),\n",
    "        dice_coef(batch_size=batch_size),\n",
    "        DiceCoefByClassAndEmptiness(cls_id=0, empty_masks_only=True, batch_size=batch_size, name='dice_coef_on_empty_cls0'),\n",
    "        DiceCoefByClassAndEmptiness(cls_id=1, empty_masks_only=True, batch_size=batch_size, name='dice_coef_on_empty_cls1'),\n",
    "        DiceCoefByClassAndEmptiness(cls_id=2, empty_masks_only=True, batch_size=batch_size, name='dice_coef_on_empty_cls2'),\n",
    "        DiceCoefByClassAndEmptiness(cls_id=3, empty_masks_only=True, batch_size=batch_size, name='dice_coef_on_empty_cls3'),\n",
    "        DiceCoefByClassAndEmptiness(cls_id=0, empty_masks_only=False, batch_size=batch_size, name='dice_coef_on_non_empty_cls0'),\n",
    "        DiceCoefByClassAndEmptiness(cls_id=1, empty_masks_only=False, batch_size=batch_size, name='dice_coef_on_non_empty_cls1'),\n",
    "        DiceCoefByClassAndEmptiness(cls_id=2, empty_masks_only=False, batch_size=batch_size, name='dice_coef_on_non_empty_cls2'),\n",
    "        DiceCoefByClassAndEmptiness(cls_id=3, empty_masks_only=False, batch_size=batch_size, name='dice_coef_on_non_empty_cls3'),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Initial Weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls checkpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_checkpoint_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#date_str = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
    "# date_str = '20191005-175252'\n",
    "#date_str = '20191018-010111'\n",
    "date_str = '20191023-030656'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_name = f'{model_checkpoint_name}_{date_str}'\n",
    "checkpoint_path = f'checkpoints/{checkpoint_name}/cp-{checkpoint_name}' + '-{epoch:04d}.ckpt'\n",
    "checkpoint_dir = os.path.dirname(checkpoint_path)\n",
    "\n",
    "initial_epoch = 0\n",
    "checkpoints = os.listdir(checkpoint_dir)\n",
    "checkpoints.sort()\n",
    "if len(checkpoints) == 0:\n",
    "    print('No checkpoints found. Starting from scratch.')\n",
    "else:\n",
    "    latest_checkpoint = os.path.join(checkpoint_dir, checkpoints[-1])\n",
    "    print(f'Loading weights from {latest_checkpoint}')\n",
    "    last_epoch = latest_checkpoint.split('-')[-1]\n",
    "    last_epoch = last_epoch.split('.')[0]\n",
    "    initial_epoch = int(last_epoch)\n",
    "    seg_model.load_weights(latest_checkpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use new model name?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_checkpoint_name = 'resnet34_unet_imgaug_focal_red_lr'\n",
    "date_str = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
    "#initial_epoch = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_name = f'{model_checkpoint_name}_{date_str}'\n",
    "checkpoint_path = f'checkpoints/{checkpoint_name}/cp-{checkpoint_name}' + '-{epoch:04d}.ckpt'\n",
    "checkpoint_dir = os.path.dirname(checkpoint_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(checkpoint_dir, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create checkpoint callback\n",
    "checkpoint_cb = tf.keras.callbacks.ModelCheckpoint(\n",
    "    checkpoint_path,\n",
    "    monitor='loss', #'val__dice_coef',#'val_loss',\n",
    "    save_best_only=True,\n",
    "    mode='auto',#'max',#'auto',\n",
    "    save_weights_only=True,\n",
    "    verbose=1)\n",
    "\n",
    "logdir = f'logs/{checkpoint_name}-{initial_epoch}'\n",
    "callbacks = [\n",
    "    tf.keras.callbacks.TensorBoard(log_dir=logdir),\n",
    "    checkpoint_cb,\n",
    "]\n",
    "\n",
    "results = seg_model.fit_generator(\n",
    "    train_data,\n",
    "    epochs=400,\n",
    "    verbose=2,\n",
    "    callbacks=callbacks,\n",
    "    validation_data=val_data,\n",
    "    steps_per_epoch=len(train_data),\n",
    "    validation_steps=len(val_data),\n",
    "    validation_freq=1,#3,\n",
    "    initial_epoch=initial_epoch,\n",
    "    workers=6,\n",
    "    use_multiprocessing=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tune Thresholds on Validation Set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def per_class_binary_cross_entropy(y_pred, y_true):\n",
    "    assert len(y_true.shape) == 3\n",
    "    eps = 0.000001\n",
    "    y_pred = np.clip(y_pred, eps, 1 - eps)\n",
    "    per_pixel_class_cross_entropy = \\\n",
    "        y_true * np.log(y_pred) + (1 - y_true) * np.log(1 - y_pred)\n",
    "    per_class_cross_entropy = np.mean(per_pixel_class_cross_entropy, axis=(0, 1))\n",
    "    return per_class_cross_entropy\n",
    "\n",
    "def per_class_iou_score(y_pred, y_true):\n",
    "    assert len(y_true.shape) == 3\n",
    "    intersection = np.sum(y_pred * y_true, axis=(0, 1))\n",
    "    union = np.sum(y_pred, axis=(0, 1)) + np.sum(y_true, axis=(0, 1)) - intersection\n",
    "    \n",
    "    iou = []\n",
    "    for i in range(intersection.shape[0]):\n",
    "        if union[i] == 0:\n",
    "            iou.append(1.0) # Both y_pred and y_true were empty\n",
    "        else:\n",
    "            iou.append(intersection[i] / union[i])\n",
    "    return np.array(iou)\n",
    "\n",
    "def per_class_mask_not_empty(y_true):\n",
    "    return np.sum(y_true, axis=(0, 1)) > 0.5\n",
    "\n",
    "def apply_tta(img):\n",
    "    h, w, c = img.shape # Assert only 3 dimensions\n",
    "    \n",
    "    img_flip_h = img[:, ::-1, :]\n",
    "    img_flip_v = img[::-1, :, :]\n",
    "    img_flip_hv = img[::-1, ::-1, :]\n",
    "    \n",
    "    tta_batch = np.stack([img, img_flip_h, img_flip_v, img_flip_hv])\n",
    "    return tta_batch\n",
    "\n",
    "def combine_tta_preds(y_tta):\n",
    "    y_0 = y_tta[0, :, :, :]\n",
    "    y_1 = y_tta[1, :, ::-1, :]\n",
    "    y_2 = y_tta[2, ::-1, :, :]\n",
    "    y_3 = y_tta[3, ::-1, ::-1, :]\n",
    "    \n",
    "    y = np.stack([y_0, y_1, y_2, y_3])\n",
    "    y = np.mean(y, axis=0, keepdims=True)\n",
    "    return y\n",
    "\n",
    "def eval_segmentation(\n",
    "    model,\n",
    "    dataset,\n",
    "    img_list,\n",
    "    thresholds=None,\n",
    "    num_classes=4):\n",
    "    \n",
    "    binary_cross_entropy = np.zeros((len(img_list), num_classes), dtype=np.float32)\n",
    "    dice_coeff = np.zeros((len(img_list), num_classes), dtype=np.float32)\n",
    "    iou_score = np.zeros((len(img_list), num_classes), dtype=np.float32)\n",
    "    mask_not_empty = np.zeros((len(img_list), num_classes), dtype=np.bool)\n",
    "    \n",
    "    if thresholds is None:\n",
    "        thresholds = [0.5, 0.5, 0.5, 0.5]\n",
    "    thresholds = np.array(thresholds)\n",
    "\n",
    "    for i, img_name in enumerate(img_list):\n",
    "        img, ann = dataset.get_example_from_img_name(img_name)\n",
    "        \n",
    "        #img_batch = np.expand_dims(img, axis=0)\n",
    "        img_batch = apply_tta(img)\n",
    "        y = model.predict(img_batch)\n",
    "        y = combine_tta_preds(y)\n",
    "        \n",
    "        # Binarize predictions\n",
    "        y_bin = np.zeros_like(y, dtype=np.uint8)\n",
    "        y_bin[y > thresholds] = 1\n",
    "\n",
    "        binary_cross_entropy[i, :] = per_class_binary_cross_entropy(y[0, :, :, :], ann)\n",
    "        dice_coeff[i, :] = per_class_dice_coeff(y_bin[0, :, :, :], ann)\n",
    "        iou_score[i, :] = per_class_iou_score(y_bin[0, :, :, :], ann)\n",
    "        mask_not_empty[i, :] = per_class_mask_not_empty(ann)\n",
    "    \n",
    "    print(f'Mean dice coeff: {np.mean(dice_coeff)}')\n",
    "    print(f'Mean binary cross entropy: {np.mean(binary_cross_entropy)}')\n",
    "    print(f'Mean IoU: {np.mean(iou_score)}')\n",
    "    \n",
    "    for i in range(cfg['NUM_CLASSES']):\n",
    "        print('*******************')\n",
    "        print(f'***** Class {i} *****')\n",
    "        print('*******************')\n",
    "        print(f'Mean dice coeff: {np.mean(dice_coeff[:, i])}')\n",
    "        print(f'Mean dice coeff (with mask): {np.mean(dice_coeff[:, i][mask_not_empty[:, i]])}')\n",
    "        print(f'Mean dice coeff (no mask): {np.mean(dice_coeff[:, i][~mask_not_empty[:, i]])}')\n",
    "        print(f'Mean IoU: {np.mean(iou_score[:, i])}')\n",
    "        print(f'Mean IoU (with mask): {np.mean(iou_score[:, i][mask_not_empty[:, i]])}')\n",
    "        print(f'Mean IoU (no mask): {np.mean(iou_score[:, i][~mask_not_empty[:, i]])}')\n",
    "        print('')\n",
    "        \n",
    "    return binary_cross_entropy, dice_coeff, iou_score, mask_not_empty"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for thresholds in [[0.3, 0.3, 0.3, 0.3], [0.5, 0.5, 0.5, 0.5], [0.7, 0.7, 0.7, 0.7]]:#, [0.6, 0.6, 0.6, 0.6], [0.7, 0.7, 0.7, 0.7]]:\n",
    "    print(f'\\n\\nThresholds: {thresholds}')\n",
    "    binary_cross_entropy, dice_coeff, iou_score, mask_not_empty = eval_segmentation(\n",
    "        seg_model,\n",
    "        val_data,\n",
    "        val_imgs,\n",
    "        thresholds=thresholds,\n",
    "        num_classes=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluate on Test Set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "binary_cross_entropy, dice_coeff, iou_score, mask_not_empty = eval_segmentation(\n",
    "    seg_model,\n",
    "    test_data,\n",
    "    test_imgs,\n",
    "    thresholds=[0.5, 0.5, 0.5, 0.5],\n",
    "    num_classes=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class_id = None\n",
    "mask_only = True\n",
    "\n",
    "scores = None\n",
    "if class_id is None:\n",
    "    print('Worst scores for all classes:')\n",
    "    scores = np.mean(dice_coeff, axis=-1)\n",
    "else:\n",
    "    print(f'Worst scores for class {class_id}')\n",
    "    scores = dice_coeff[:, class_id]\n",
    "\n",
    "indices = np.argsort(scores) # Indices of worst images\n",
    "\n",
    "if mask_only and class_id is not None:\n",
    "    print('Including scores for non-empty ground truth masks only.')\n",
    "    mask_only_indices = np.where(mask_not_empty[:, class_id])\n",
    "    mask_only_indices = set(mask_only_indices[0].tolist())\n",
    "    indices = [index for index in indices if index in mask_only_indices]\n",
    "\n",
    "for i in indices:\n",
    "    print(f'{i}: {scores[i]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize Image Prediction\n",
    "img_id = 145\n",
    "thresh = [0.5, 0.5, 0.5, 0.5]\n",
    "\n",
    "img_name = test_imgs[img_id]\n",
    "img, ann = test_data.get_example_from_img_name(img_name)\n",
    "img_batch = np.expand_dims(img, axis=0)\n",
    "y = seg_model.predict(img_batch)\n",
    "plt.figure(figsize=(10, 3))\n",
    "plt.imshow(visualize_segmentations(np.repeat(img, 3, axis=-1), ann))\n",
    "plt.show()\n",
    "\n",
    "for i in range(y.shape[-1]):\n",
    "    plt.figure(figsize=(12.5, 3))\n",
    "    plt.imshow(y[0, :, :, i])\n",
    "    plt.colorbar()\n",
    "\n",
    "for i in range(y.shape[-1]):\n",
    "    plt.figure(figsize=(10, 3))\n",
    "    plt.imshow(y[0, :, :, i] > thresh[i])\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Save HDF5 Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "date_str = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
    "seg_model.save(f'resnet_seg_model_imgaug_{date_str}.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "date_str = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
    "tf.keras.models.save_model(\n",
    "    seg_model,\n",
    "    f'tf_resnet_seg_model_{date_str}.h5',\n",
    "    include_optimizer=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}