modeltrain_autoencoder.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "https://github.com/karolzak/keras-unet/blob/master/notebooks/kz-whale-tails.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import glob\n",
    "import os\n",
    "import sys\n",
    "from pathlib import Path\n",
    "from PIL import Image\n",
    "\n",
    "import tensorflow as tf\n",
    "import tensorflow.keras as keras\n",
    "from tensorflow.keras.callbacks import ModelCheckpoint\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "DATA_DIR = Path(\"./data/tiles/ls78_ae\")\n",
    "BATCH_SIZE = 2\n",
    "\n",
    "MAX_X = 255\n",
    "MAX_Y = 65535  # TODO: Tilegen outputs rgb image, not grayscale\n",
    "\n",
    "BAND_DIRS = sorted(list(DATA_DIR.glob(\"B*\"))) # [::-1]  # reverse for rgb\n",
    "BASE_DIR = DATA_DIR / \"B1\"\n",
    "\n",
    "display(BAND_DIRS)\n",
    "display(list())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_paths = list(BASE_DIR.glob(\"*.png\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_fname(fname):\n",
    "    # div by maxy, but mul by maxx in anticipation of preprocess_transform\n",
    "    bands = np.asarray([np.array(Image.open(band_dir / fname)) for band_dir in BAND_DIRS])\n",
    "#     bands = bands[1:4] # [::-1]\n",
    "    mchannel = np.dstack(bands).astype('uint8')\n",
    "    return mchannel\n",
    "\n",
    "def preprocess_transform(ds):\n",
    "    return (ds / MAX_X).astype('float32')\n",
    "\n",
    "def rgb_transform(ds):\n",
    "#     return ds # bands 4, 3, 2, -> rgb\n",
    "#     return np.flip(ds, 3)  # Bands 2, 3, 4 -> rgb\n",
    "    return np.flip(ds[:,:,:,1:4], 3)  # Bands 1, 2, 3, 4, 5, 7 -> rgb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(base_paths[0].name)\n",
    "\n",
    "x0 = read_fname(base_paths[0].name)\n",
    "\n",
    "display(x0.shape)\n",
    "# display(y0.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "img_lst = []\n",
    "\n",
    "for base_path in base_paths:\n",
    "    img = read_fname(base_path.name)\n",
    "    img_lst.append(img)\n",
    "\n",
    "img_np = np.asarray(img_lst)\n",
    "\n",
    "del img_lst\n",
    "\n",
    "display(f'{img_np.shape=}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from keras_unet.utils import plot_imgs, plot_aeimgs\n",
    "\n",
    "display(img_np.min())\n",
    "display(img_np.max())\n",
    "\n",
    "plot_aeimgs(org_imgs=rgb_transform(img_np) / 255, nm_img_to_plot=5, figsize=6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "x_train, x_val, _, _ = train_test_split(img_np, img_np, test_size=0.30, random_state=0)\n",
    "\n",
    "print(\"x_train: \", x_train.shape)\n",
    "# print(\"y_train: \", y_train.shape)\n",
    "print(\"x_val: \", x_val.shape)\n",
    "# print(\"y_val: \", y_val.shape)\n",
    "\n",
    "input_shape = x_train[0].shape\n",
    "display(f'{input_shape=}')\n",
    "\n",
    "train_len = len(x_train)\n",
    "validation_len = len(x_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_val = preprocess_transform(x_val)\n",
    "\n",
    "x_val_tf = tf.convert_to_tensor(x_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from keras_unet.utils import get_augmented\n",
    "\n",
    "train_gen = get_augmented(\n",
    "    x_train, x_train, \n",
    "#     x_val_tf, y_val_tf,\n",
    "    batch_size=BATCH_SIZE,\n",
    "    data_gen_args = dict(\n",
    "#         rescale=1 / 255,\n",
    "        preprocessing_function=preprocess_transform,\n",
    "        rotation_range=360.,\n",
    "        width_shift_range=0.05,\n",
    "        height_shift_range=0.05,\n",
    "        shear_range=40,\n",
    "        zoom_range=0.2,\n",
    "#         brightness_range=[0.6,1.4],  # only works on 1 or 3 channel images\n",
    "        channel_shift_range=10.0,\n",
    "        horizontal_flip=True,\n",
    "        vertical_flip=True,\n",
    "        fill_mode='reflect'\n",
    "    ))\n",
    "\n",
    "# train_gen = get_augmented(\n",
    "#     x_train, y_train, batch_size=2,\n",
    "#     data_gen_args=dict(\n",
    "#         rotation_range=0.,\n",
    "#         width_shift_range=0.00,\n",
    "#         height_shift_range=0.00,\n",
    "#         shear_range=0,\n",
    "#         zoom_range=0.0,\n",
    "#         horizontal_flip=True,\n",
    "#         vertical_flip=True,\n",
    "#         fill_mode='constant'\n",
    "#     )\n",
    "# )\n",
    "\n",
    "del x_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "sample_batch = next(train_gen)\n",
    "\n",
    "xx, yy = sample_batch\n",
    "print(xx.shape, yy.shape)\n",
    "from keras_unet.utils import plot_imgs\n",
    "\n",
    "plot_aeimgs(org_imgs=rgb_transform(xx), nm_img_to_plot=2, figsize=6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from bcdunet import BCDU_net_D3, BCDU_net_D1\n",
    "\n",
    "# input_shape = (512, 512, 3)\n",
    "# model = BCDU_net_D3(input_size=input_shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from keras_unet.models import custom_unet\n",
    "\n",
    "# # input_shape = (512, 512, 6)\n",
    "\n",
    "def create_model_unet():\n",
    "    model = custom_unet(\n",
    "        input_shape,\n",
    "        filters=32,\n",
    "        use_batch_norm=True,\n",
    "        dropout=0.5,  # 0.3\n",
    "        dropout_change_per_layer=0.0,\n",
    "        use_attention=False,\n",
    "        use_squeeze_on_enc=True,\n",
    "        use_squeeze_on_dec=True,\n",
    "        num_layers=5,\n",
    "        skip_connections=False,  # auto encoder\n",
    "        num_classes=6,  # 6 output channels\n",
    "    )\n",
    "    return model\n",
    "\n",
    "model = create_model_unet()\n",
    "\n",
    "# model = custom_unet(\n",
    "#     input_shape,\n",
    "#     filters=32,\n",
    "#     use_batch_norm=True,\n",
    "#     dropout=0.5,  # 0.3\n",
    "#     dropout_change_per_layer=0.0,\n",
    "#     use_attention=False,\n",
    "#     use_squeeze_on_enc=True,\n",
    "#     use_squeeze_on_dec=True,\n",
    "#     num_layers=5,\n",
    "# )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "tf.keras.utils.plot_model(model)\n",
    "\n",
    "display(model.summary())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_filename = 'model_trainchpt.h5'\n",
    "callback_checkpoint = ModelCheckpoint(\n",
    "    model_filename, \n",
    "    verbose=1, \n",
    "    monitor='val_loss', \n",
    "    save_best_only=True,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://github.com/ailias/Focal-Loss-implement-on-Tensorflow/blob/master/focal_loss.py\n",
    "\n",
    "from tensorflow.python.ops import array_ops\n",
    "\n",
    "def focal_loss(prediction_tensor, target_tensor, weights=None, alpha=0.25, gamma=2):\n",
    "    r\"\"\"Compute focal loss for predictions.\n",
    "        Multi-labels Focal loss formula:\n",
    "            FL = -alpha * (z-p)^gamma * log(p) -(1-alpha) * p^gamma * log(1-p)\n",
    "                 ,which alpha = 0.25, gamma = 2, p = sigmoid(x), z = target_tensor.\n",
    "    Args:\n",
    "     prediction_tensor: A float tensor of shape [batch_size, num_anchors,\n",
    "        num_classes] representing the predicted logits for each class\n",
    "     target_tensor: A float tensor of shape [batch_size, num_anchors,\n",
    "        num_classes] representing one-hot encoded classification targets\n",
    "     weights: A float tensor of shape [batch_size, num_anchors]\n",
    "     alpha: A scalar tensor for focal loss alpha hyper-parameter\n",
    "     gamma: A scalar tensor for focal loss gamma hyper-parameter\n",
    "    Returns:\n",
    "        loss: A (scalar) tensor representing the value of the loss function\n",
    "    \"\"\"\n",
    "    sigmoid_p = tf.nn.sigmoid(prediction_tensor)\n",
    "    zeros = array_ops.zeros_like(sigmoid_p, dtype=sigmoid_p.dtype)\n",
    "    \n",
    "    # For poitive prediction, only need consider front part loss, back part is 0;\n",
    "    # target_tensor > zeros <=> z=1, so poitive coefficient = z - p.\n",
    "    pos_p_sub = array_ops.where(target_tensor > zeros, target_tensor - sigmoid_p, zeros)\n",
    "    \n",
    "    # For negative prediction, only need consider back part loss, front part is 0;\n",
    "    # target_tensor > zeros <=> z=1, so negative coefficient = 0.\n",
    "    neg_p_sub = array_ops.where(target_tensor > zeros, zeros, sigmoid_p)\n",
    "    per_entry_cross_ent = - alpha * (pos_p_sub ** gamma) * tf.math.log(tf.clip_by_value(sigmoid_p, 1e-8, 1.0)) \\\n",
    "                          - (1 - alpha) * (neg_p_sub ** gamma) * tf.math.log(tf.clip_by_value(1.0 - sigmoid_p, 1e-8, 1.0))\n",
    "    return tf.reduce_sum(per_entry_cross_ent)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.optimizers import Adam, SGD\n",
    "from keras_unet.metrics import iou, iou_thresholded\n",
    "from keras_unet.losses import jaccard_distance\n",
    "\n",
    "def closs(y_true, y_pred):\n",
    "    return 2 * jaccard_distance(y_true, y_pred) + tf.keras.losses.binary_crossentropy(y_true, y_pred)\n",
    "\n",
    "model.compile(\n",
    "    optimizer=Adam(),\n",
    "#     optimizer=\"adam\",\n",
    "    #optimizer=SGD(lr=0.01, momentum=0.99),\n",
    "    loss='binary_crossentropy',\n",
    "#     loss=jaccard_distance,\n",
    "#     loss=focal_loss,\n",
    "#     loss=closs,\n",
    "    metrics=[iou, iou_thresholded],\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.load_weights('autoencoder.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "history = model.fit(\n",
    "    train_gen,\n",
    "    steps_per_epoch=train_len // BATCH_SIZE,\n",
    "    epochs=500,\n",
    "    validation_data=(x_val_tf, x_val_tf),\n",
    "    validation_batch_size=BATCH_SIZE,\n",
    "    callbacks=[callback_checkpoint]\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from keras_unet.utils import plot_segm_history\n",
    "\n",
    "plot_segm_history(history)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_weights('autoencoder.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.load_weights(model_filename)\n",
    "# model.load_weights('model_trainchpt.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_off = 1000\n",
    "y_pred = model.predict(x_val[pred_off:pred_off + 10], batch_size=BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from keras_unet.utils import plot_imgs, plot_aeimgs\n",
    "\n",
    "display(len(x_val))\n",
    "plot_aeimgs(\n",
    "    org_imgs=rgb_transform(x_val[pred_off:pred_off + 10]),\n",
    "    pred_imgs=rgb_transform(y_pred),\n",
    "    nm_img_to_plot=10, figsize=4)\n"
   ]
  },
  {
   "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}