{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Salt Body Interpretation on Seismic Using Sagemaker and MXNET"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This notebook contains a tutorial on how to build a deep learning (semantic segmentation) model for automatic salt interpretation. \n",
"* The fully convolutional architecture known as UNet for semantic segmentation\n",
"* How to train UNet in Amazon SageMaker, and deploy to an inference endpoint\n",
"\n",
"Import the following modules:"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## IMPORTANT"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Depending on when you started your Sagemaker notebook, you might need to re-install \"scikit-image\", \"scikit-learn\", \"numpy\" and \"scipy\" libraries as some functions used in this notebook may not be available in other versions. If you get errors about using any of these 4 libraries, you can un-comment below 4 lines of code and run it."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# !pip install scikit-learn==0.16.0\n",
"# !pip install scikit-image==0.12.2\n",
"# !pip install scipy==1.2.1 \n",
"# !pip install numpy==1.16.4 "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": false
},
"outputs": [],
"source": [
"import mxnet as mx\n",
"import segmentation_methods as sgm\n",
"from mxnet import ndarray as F\n",
"import numpy as np\n",
"import urllib\n",
"from PIL import Image\n",
"np.random.seed(1984)\n",
"import glob\n",
"import os\n",
"import urllib\n",
"import zipfile\n",
"from scipy.misc import imresize\n",
"from sklearn.cross_validation import train_test_split\n",
"import scipy.io as sio\n",
"from skimage import measure\n",
"import time\n",
"import matplotlib.pyplot as plt\n",
"import warnings\n",
"warnings.filterwarnings('ignore')\n",
"%matplotlib inline"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Data ingestion"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The dataset use din this study is provided by TGS. Here is a link to data: https://www.kaggle.com/c/tgs-salt-identification-challenge\n",
"We load the names of the files containing the labels. For each image that has a label, we load that name."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"image_dir = 'data/train/images/'\n",
"image_files = sgm.get_file_path_list(image_dir)\n",
"label_dir = 'data/train/masks/'\n",
"label_files = sgm.get_file_path_list(label_dir)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Here we load the masks, and convert them into a binary format\"\n",
"* If the images are not the same resolution, so they are resized to a constant 820x550. Any interpolated label values greater than zero are set to one.\n",
"* MXNet requires the input to have dimension (batch, channel, height, width), so these alterations are made."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"image1 = np.array(Image.open(image_files[100]).resize((550,820)))\n",
"mask1 = np.array(Image.open(label_files[100]).resize((550,820)))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"plt.figure(figsize=(12,12))\n",
"plt.subplot(131)\n",
"plt.title('Image')\n",
"plt.imshow(image1)\n",
"\n",
"plt.subplot(132)\n",
"plt.imshow(mask1)\n",
"plt.title('Salt Mask')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Stack images"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": false
},
"outputs": [],
"source": [
"X = []\n",
"Y = []\n",
"for i in range(len(image_files)):\n",
" mask = (Image.open(label_files[i]))\n",
" mask = (imresize(mask, (820, 550)) > 0).astype(np.uint8) # interpolate to 820 x 550\n",
" image = np.array(Image.open(image_files[i]).resize((550,820)))\n",
" X.append(image)\n",
" Y.append(mask)\n",
"X = np.transpose(np.stack(X), axes=(0, 3, 1, 2))\n",
"Y = np.expand_dims(np.stack(Y), 1)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We have 4000 observations."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"X.shape"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Data splitting and augmentation (random cropping)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We're going to generate a new data-set through random cropping of our existing images. Before we do that, we need to split the data into training and validation data (if we did crops first, and then split, we run the risk of data leakage)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"train_X, validation_X, train_Y, validation_Y = train_test_split(X, Y, test_size=0.2, random_state=1984)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"train_X_boxes, train_Y_boxes = sgm.extract_boxes(train_X, train_Y)\n",
"validation_X_boxes, validation_Y_boxes = sgm.extract_boxes(validation_X, validation_Y)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Next, we generate the random crops."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"train_X_crops, train_Y_crops = sgm.generate_random_crops(train_X_boxes, train_Y_boxes, num_patches=3)\n",
"validation_X_crops, validation_Y_crops = sgm.generate_random_crops(validation_X_boxes, validation_Y_boxes, num_patches=3)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"train_X_crops.shape"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Next, we'll save the generated data locally so we can avoid generating again."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"if not os.path.exists('/dev/shm/salt/segmentation_data'): \n",
" os.mkdir('/dev/shm/salt/')\n",
" os.mkdir('/dev/shm/salt/segmentation_data')\n",
"np.save('/dev/shm/salt/segmentation_data/train_X_crops.npy', train_X_crops)\n",
"np.save('/dev/shm/salt/segmentation_data/train_Y_crops.npy', train_Y_crops)\n",
"np.save('/dev/shm/salt/segmentation_data/validation_X_crops.npy', validation_X_crops)\n",
"np.save('/dev/shm/salt/segmentation_data/validation_Y_crops.npy', validation_Y_crops)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"train_X_crops = np.load('/dev/shm/salt/segmentation_data/train_X_crops.npy')\n",
"train_Y_crops = np.load('/dev/shm/salt/segmentation_data/train_Y_crops.npy')\n",
"validation_X_crops = np.load('/dev/shm/salt/segmentation_data/validation_X_crops.npy')\n",
"validation_Y_crops = np.load('/dev/shm/salt/segmentation_data/validation_Y_crops.npy')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# SageMaker"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We're going to proceed by defining the UNet Network for binary segmentation using Sagemaker."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now we are ready to define a training job in SageMaker to do training at scale."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import boto3\n",
"import sagemaker\n",
"from sagemaker.mxnet import MXNet\n",
"from sagemaker import get_execution_role"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"role = get_execution_role()\n",
"sagemaker_session = sagemaker.Session()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We need to upload the data to S3 so the instances launched for the training job can pull the data down."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"inputs = sagemaker_session.upload_data(path='/dev/shm/salt/segmentation_data', key_prefix='sagemaker_data')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Finally, we create the MXNet SageMaker model estimator using the **Bring your own script** paradigm. We've defined a script, segmentation.py, that runs the training loop for UNet in MXNet Symbolic. To do this, we follow the conventions defined for the SageMaker Python SDK [here](https://github.com/aws/sagemaker-python-sdk)."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Train"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"#SAGEMAKER NOTEBOOK CODE\n",
"sagemaker_net = MXNet(\"segmentation.py\", \n",
" role=role, \n",
" train_instance_count=2, \n",
" train_instance_type=\"ml.p3.16xlarge\",\n",
" sagemaker_session=sagemaker_session,\n",
" framework_version=\"1.2\",\n",
" hyperparameters={\n",
" 'data_shape': (3, 256, 256),\n",
" 'batch_size': 64, \n",
" 'epochs': 100, \n",
" 'learning_rate': 1E-3, \n",
" 'num_gpus': 1,\n",
" })\n",
"\n",
"sagemaker_net.fit(inputs)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Once the training is complete, we can launch an endpoint server that serves inference with our trained model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"sagemaker_predictor = sagemaker_net.deploy(initial_instance_count=1, instance_type='ml.p2.xlarge')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Test"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"test_iter = mx.io.NDArrayIter(data = validation_X_crops, label=validation_Y_crops, batch_size=1, shuffle=False)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can send test data to the inference endpoint:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"batch = next(test_iter)\n",
"data = batch.data[0]\n",
"label = batch.label[0]\n",
"response = sagemaker_predictor.predict(data.asnumpy().tolist())\n",
"output = np.array(response[0])"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def post_process_mask(label, p=0.5):\n",
" return (np.where(label > p, 1, 0)).astype('uint8')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"width = 12\n",
"height = 12\n",
"plt.figure(figsize=(width, height))\n",
"plt.subplot(331)\n",
"plt.title('Input')\n",
"plt.imshow(np.transpose(data.asnumpy()[0], (1,2,0)).astype(np.uint8))\n",
"plt.subplot(332)\n",
"plt.title('Prediction')\n",
"plt.imshow(post_process_mask(output[0]), cmap=plt.cm.gray)\n",
"plt.subplot(333)\n",
"plt.title('Mask')\n",
"plt.imshow(label[0][0].asnumpy(), cmap=plt.cm.gray)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Don't forget to delete your endpoint when you're done with it."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"sagemaker_net.delete_endpoint()"
]
}
],
"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.7.3"
}
},
"nbformat": 4,
"nbformat_minor": 2
}