{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom image classification example\n",
    "\n",
    "This notebook is based off the [example created by Gabe Hollombe](https://github.com/gabehollombe-aws/jupyter-notebooks). It has been simplified and modified for a defect detection example project. If interested for a more in-depth explanation of the image classification process, please visit the original notebook. \n",
    "Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved. SPDX-License-Identifier: MIT-0 \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1 - Specify training data location\n",
    "\n",
    "Provide a bucket containing the image dataset. Organize images into folders by class. \n",
    "\n",
    "```\n",
    "images_to_classify\n",
    "├── defect_free\n",
    "│   ├── 1.jpg\n",
    "│   ├── 2.jpg\n",
    "|   ├── 3.jpg\n",
    "│   └── . . .\n",
    "└── defective\n",
    "│   ├── 1.jpg\n",
    "│   ├── 2.jpg    \n",
    "│   ├── 3.jpg\n",
    "│   ├── . . .\n",
    "└── . . .\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "# The S3 bucket name where the image dataset is located\n",
    "data_bucket_name='SomeBucket'\n",
    "\n",
    "# The name of the folder containing the dataset\n",
    "dataset_name = 'images_to_classify'\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2 - Setup environment\n",
    "\n",
    "The following imports the neccesary python libraries and fetches the execution role for the notebook instance. It also fetches the built-in image-classification algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "from sagemaker import get_execution_role\n",
    "from sagemaker.amazon.amazon_estimator import get_image_uri\n",
    "\n",
    "role = get_execution_role()\n",
    "sess = sagemaker.Session()\n",
    "\n",
    "training_image = get_image_uri(sess.boto_region_name, 'image-classification', repo_version=\"latest\")\n",
    "\n",
    "# Find im2rec in our environment and set up some other vars in our environemnt\n",
    "\n",
    "base_dir='/tmp'\n",
    "\n",
    "%env BASE_DIR=$base_dir\n",
    "%env S3_DATA_BUCKET_NAME = $data_bucket_name\n",
    "%env DATASET_NAME = $dataset_name\n",
    "\n",
    "import sys,os\n",
    "\n",
    "suffix='/mxnet/tools/im2rec.py'\n",
    "im2rec = list(filter( (lambda x: os.path.isfile(x + suffix )), sys.path))[0] + suffix\n",
    "%env IM2REC=$im2rec\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3 - Download image dataset from Amazon S3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pull our images from S3\n",
    "!aws s3 sync s3://$S3_DATA_BUCKET_NAME/$DATASET_NAME $BASE_DIR/$DATASET_NAME "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4 - Create RecordIO files from images\n",
    "\n",
    "The images need to be converted into RecordIO files. Using im2rec.py, These files contain binary data of the images indexed by class, which is infered by the folder structure. Two files are made for training and validation. \n",
    "\n",
    "For more information see https://docs.aws.amazon.com/sagemaker/latest/dg/image-classification.html#IC-inputoutput"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "# Use the IM2REC script to convert our images into RecordIO files\n",
    "\n",
    "# Clean up our working dir of existing LST and REC files\n",
    "cd $BASE_DIR\n",
    "rm *.rec\n",
    "rm *.lst\n",
    "\n",
    "# First we need to create two LST files (training and test lists), noting the correct label class for each image\n",
    "# We'll also save the output of the LST files command, since it includes a list of all of our label classes\n",
    "echo \"Creating LST files\"\n",
    "python $IM2REC --list --recursive --pass-through --test-ratio=0.3 --train-ratio=0.7 $DATASET_NAME $DATASET_NAME > ${DATASET_NAME}_classes\n",
    "\n",
    "echo \"Label classes:\"\n",
    "cat ${DATASET_NAME}_classes\n",
    "\n",
    "# Then we create RecordIO files from the LST files\n",
    "echo \"Creating RecordIO files\"\n",
    "python $IM2REC --num-thread=4 ${DATASET_NAME}_train.lst $DATASET_NAME\n",
    "python $IM2REC --num-thread=4 ${DATASET_NAME}_test.lst $DATASET_NAME\n",
    "ls -lh *.rec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5 - Upload RecordIO files to S3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Upload our train and test RecordIO files to S3 in the bucket that our sagemaker session is using\n",
    "bucket = sess.default_bucket()\n",
    "\n",
    "s3train_path = 's3://{}/{}/train/'.format(bucket, dataset_name)\n",
    "s3validation_path = 's3://{}/{}/validation/'.format(bucket, dataset_name)\n",
    "\n",
    "# Clean up any existing data\n",
    "!aws s3 rm s3://{bucket}/{dataset_name}/train --recursive\n",
    "!aws s3 rm s3://{bucket}/{dataset_name}/validation --recursive\n",
    "\n",
    "# Upload the rec files to the train and validation channels\n",
    "!aws s3 cp /tmp/{dataset_name}_train.rec $s3train_path\n",
    "!aws s3 cp /tmp/{dataset_name}_test.rec $s3validation_path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 6 - Configure training model\n",
    "\n",
    "For this example nothing needs to be changed. Batch size can be modified to reflect your sample size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# The minimum batch size should be less than the total sample size for each class. \n",
    "batch_size = 2\n",
    "\n",
    "train_data = sagemaker.session.s3_input(\n",
    "    s3train_path, \n",
    "    distribution='FullyReplicated', \n",
    "    content_type='application/x-recordio', \n",
    "    s3_data_type='S3Prefix'\n",
    ")\n",
    "\n",
    "validation_data = sagemaker.session.s3_input(\n",
    "    s3validation_path, \n",
    "    distribution='FullyReplicated', \n",
    "    content_type='application/x-recordio', \n",
    "    s3_data_type='S3Prefix'\n",
    ")\n",
    "\n",
    "data_channels = {'train': train_data, 'validation': validation_data}\n",
    "\n",
    "s3_output_location = 's3://{}/{}/output'.format(bucket, dataset_name)\n",
    "\n",
    "image_classifier = sagemaker.estimator.Estimator(\n",
    "    training_image,\n",
    "    role, \n",
    "    train_instance_count=1, \n",
    "    train_instance_type='ml.p3.2xlarge',\n",
    "    output_path=s3_output_location,\n",
    "    sagemaker_session=sess\n",
    ")\n",
    "\n",
    "num_classes=! ls -l {base_dir}/{dataset_name} | wc -l\n",
    "num_classes=int(num_classes[0]) - 1\n",
    "\n",
    "num_training_samples=! cat {base_dir}/{dataset_name}_train.lst | wc -l\n",
    "num_training_samples = int(num_training_samples[0])\n",
    "\n",
    "# Learn more about the Sagemaker built-in Image Classifier hyperparameters here: https://docs.aws.amazon.com/sagemaker/latest/dg/IC-Hyperparameter.html\n",
    "\n",
    "# These hyperparameters we won't want to change, as they define things like\n",
    "# the size of the images we'll be sending for input, the number of training classes we have, etc.\n",
    "base_hyperparameters=dict(\n",
    "    use_pretrained_model=1,\n",
    "    image_shape='3,224,224',\n",
    "    num_classes=num_classes,\n",
    "    num_training_samples=num_training_samples,\n",
    ")\n",
    "\n",
    "# These are hyperparameters we may want to tune, as they can affect the model training success:\n",
    "hyperparameters={\n",
    "    **base_hyperparameters, \n",
    "    **dict(\n",
    "        learning_rate=0.001,\n",
    "        mini_batch_size= batch_size,\n",
    "    )\n",
    "}\n",
    "\n",
    "\n",
    "image_classifier.set_hyperparameters(**hyperparameters)\n",
    "\n",
    "hyperparameters\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 7 - Start training job\n",
    "\n",
    "This will dispatch a job to Amazon SageMaker to begin a training job. Upon completion it will upload a training model to Amazon S3 in a generate bucket for SageMaker."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "import time\n",
    "now = str(int(time.time()))\n",
    "training_job_name = 'IC-' + dataset_name.replace('_', '-') + '-' + now\n",
    "\n",
    "image_classifier.fit(inputs=data_channels, job_name=training_job_name, logs=True)\n",
    "\n",
    "job = image_classifier.latest_training_job\n",
    "model_path = f\"{base_dir}/{job.name}\"\n",
    "\n",
    "print(f\"\\n\\n Finished training! The model is available for download at: {image_classifier.output_path}/{job.name}/output/model.tar.gz\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 8 - Deploy Model Endpoint\n",
    "\n",
    "Deploys a REST endpoint that can be invoked in order to run images against for classification."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "%%time\n",
    "# Deploying a model to an endpoint takes a few minutes to complete\n",
    "\n",
    "deployed_endpoint = image_classifier.deploy(\n",
    "    initial_instance_count = 1,\n",
    "    instance_type = 'ml.t2.medium'\n",
    ")\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_mxnet_p36",
   "language": "python",
   "name": "conda_mxnet_p36"
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}