image_set_preparation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "# Preparation of a labeled set of aerial images from public data sources\n",
    "\n",
    "This notebook illustrates how a set of aerial images, labeled by land use classification, was generated from freely available U.S. datasets. The image sets generated in this notebook were used to train a DNN for image classification and evaluate its performance. For more detail, please see the rest of the [Embarrassingly Parallel Image Classification](https://github.com/Azure/Embarrassingly-Parallel-Image-Classification) repository.\n",
    "\n",
    "This walkthrough will demonstrate the technique using source data from Middlesex County, MA, the home of Microsoft's New England Research & Development (NERD) Center. The method was applied to many U.S. counties in parallel to generate the full training and validation sets.\n",
    "\n",
    "<img src=\"./img/data_overview/middlesex_ma.png\" />\n",
    "\n",
    "Note: it is not necessary to complete all steps in this tutorial in order to proceed to other tutorials in this series. For a fast start, complete the steps under the [Prepare an Azure Data Science Virtual Machine for image extraction](#dsvm) section, then proceed directly to the [Dataset preparation for deep learning](#prep) section."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "## Outline\n",
    "- [Source data](#sourcedata)\n",
    "   - [National Agriculture Imagery Program](#naip)\n",
    "   - [National Land Cover Database](#nlcd)\n",
    "- [Prepare an Azure Data Science Virtual Machine for image extraction](#dsvm)\n",
    "   - [Provision the VM](#provision)\n",
    "   - [Connect to the VM by remote desktop](#rd)\n",
    "   - [Install Python packages and supporting software](#install)\n",
    "   - [Download input data locally](#download)\n",
    "- [Produce image sets for training and evaluation](#production)\n",
    "   - [Overview of the approach](#overview)\n",
    "   - [Convert NAIP images from MrSID to GeoTIFF](#conversion)\n",
    "   - [Add tiling for fast data extraction](#tiling)\n",
    "   - [Find a lat-lon bounding box for the region](#box)\n",
    "   - [Get the approximate dimensions of the bounding box in meters](#meters)\n",
    "   - [Find tiles with consistent land use class](#consistent)\n",
    "   - [Extract and save images for tiles of interest](#extraction)\n",
    "   - [Dataset partitioning](#partition)\n",
    "- [Dataset preparation for deep learning](#prep)\n",
    "   - [Cognitive Toolkit (CNTK)](#cntk)\n",
    "   - [TensorFlow](#tf)\n",
    "- [Next Steps](#nextsteps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "<a name=\"sourcedata\"></a>\n",
    "## Source data\n",
    "\n",
    "<a name=\"naip\"></a>\n",
    "### National Agriculture Imagery Program\n",
    "\n",
    "The US [National Agriculture Imagery Program](https://www.fsa.usda.gov/programs-and-services/aerial-photography/imagery-programs/naip-imagery/index), run by the [US Department of Agriculture](https://www.usda.gov/)'s [Farm Service Agency](https://www.fsa.usda.gov/), provides images within one month of collection every 2 years, in natural (RGB) color with 1-meter ground sample distance. In this tutorial, we use Compressed County Mosaics obtained from the [Geospatial Data Gateway](https://gdg.sc.egov.usda.gov/); you can download the 2016 Middlesex County, MA directly from [our mirror](https://mawahstorage.blob.core.windows.net/aerialimageclassification/naipsample/ortho_imagery_NAIPM16_ma017_3344056_01.zip) if you prefer.\n",
    "\n",
    "A zoomed-out view of the 2016 NAIP data covering Middlesex County, MA is shown below (Cambridge and Boston at ESE edge):\n",
    "\n",
    "<img src=\"./img/data_overview/mediumnaip_white.png\" width=\"500px\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "<a name=\"nlcd\"></a>\n",
    "### National Land Cover Database\n",
    "\n",
    "The US [National Land Cover Database](https://www.mrlc.gov/nlcd2011.php) (NLCD), maintained by the [Multi-Resolution Land Characteristic Consortium](https://www.mrlc.gov/), provides land use labels at 30-meter spatial resolution for the United States. The [sixteen land use classes](https://www.mrlc.gov/nlcd11_leg.php) in the NLCD are organized hierarchically into major (Developed, Forested, Cultivated, &c.) and minor (e.g. Deciduous, Evergreen, or Mixed Forest) categories. Classifications are based largely on seasonally-collected satellite imagery, including data collected outside the visual spectrum. Because of the extensive post-processing required to prepare the NLCD dataset, new versions are published approximately every five years, often with a multi-year delay between data collection and publication. For this project, we use the [NLCD 2011 Land Cover](https://www.mrlc.gov/nlcd11_data.php) dataset.\n",
    "\n",
    "The colorized NLCD data for Middlesex County, MA and surrounding areas are shown below. Developed land is shown in shades of red, water/wetlands in blue/cyan, forested land in green, and cultivated land in magenta.\n",
    "\n",
    "<img src=\"./img/data_overview/mediumnlcd.png\" width=\"500px\"/>\n",
    "\n",
    "For more information on the NLCD, please see the following publication:\n",
    "\n",
    "Homer, C.G., Dewitz, J.A., Yang, L., Jin, S., Danielson, P., Xian, G., Coulston, J., Herold, N.D., Wickham, J.D., and Megown, K., 2015, [Completion of the 2011 National Land Cover Database for the conterminous United States - Representing a decade of land cover change information](http://bit.ly/1K7WjO3). Photogrammetric Engineering and Remote Sensing, v. 81, no. 5, p. 345-354 "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "<a name=\"dsvm\"></a>\n",
    "## Prepare an Azure Data Science Virtual Machine for image extraction\n",
    "\n",
    "For reproducibility and brevity, this walkthrough describes how to set up an [Azure Data Science Virtual Machine (DSVM)](https://azure.microsoft.com/en-us/marketplace/partners/microsoft-ads/standard-data-science-vm/) for image extraction and later deep neural network training. We recommend provisioning a Windows 2016 DSVM now so that you can progress seamlessly to the model training step afterward, but GPU compute is not necessary for the image extraction portion of this tutorial.\n",
    "\n",
    "If you prefer not to create a VM at this stage, you can likely adapt the steps below to install needed software on your local computer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "<a name=\"provision\"></a>\n",
    "### Provision the VM\n",
    "\n",
    "1. In the [Azure Portal](https://ms.portal.azure.com), begin provisioning a new Windows 2016 DSVM.\n",
    "    1. Click the \"+ New\" button at upper left to launch a search pane.\n",
    "    1. Type in \"Data Science Virtual Machine (DSVM) - Windows 2016\" and press Enter.\n",
    "    1. In the search results, choose the \"Data Science Virtual Machine (DSVM) - Windows 2016\" resource published by Microsoft.\n",
    "    1. After reading the description, press \"Create\" to begin customization.\n",
    "1. In the \"Basics\" pane:\n",
    "    1. Choose a name, username, and password for your VM.\n",
    "    1. Set the VM disk type to \"HDD\". (As of this writing, SSD disks are not available with GPU VMs.)\n",
    "    1. Create a new resource group with a unique name for your VM.\n",
    "        - Creating a separate resource group for the VM will enable to easily delete all resources associated with this tutorial when you are finished.\n",
    "    1. Set the location to \"East US\".\n",
    "        - Note that GPU VMs are not currently available in all regions.\n",
    "1. In the \"Size\" pane, choose a [virtual machine size](https://docs.microsoft.com/en-us/azure/virtual-machines/virtual-machines-windows-sizes) that includes a graphics card based on your needs.\n",
    "    - We recommend the \"NC6\" (single-GPU) virtual machine size for this tutorial. The code provided does not support distributed learning.\n",
    "1. On the \"Settings\" pane, click \"OK\" to continue without modifying the default settings.\n",
    "1. Confirm your selections on the \"Purchase\" pane, then click \"Purchase\" to provision the VM.\n",
    "\n",
    "VM deployment may take 20-30 minutes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "<a name=\"rd\"></a>\n",
    "### Connect to the VM by remote desktop\n",
    "\n",
    "After the VM deployment is finished, you can connect to the VM by remote desktop as follows:\n",
    "1. Navigate to the VM's pane in Azure Portal (e.g. by searching for the VM's name).\n",
    "1. Click \"Connect\" along the bar on top of the pane to download an RDP (remote desktop) file.\n",
    "1. Double-click the RDP file to start the connection.\n",
    "1. Supply the username and password you chose earlier when prompted.\n",
    "\n",
    "### Run this notebook using the VM's Jupyter notebook server\n",
    "\n",
    "We recommend downloading this notebook on the VM and running it with the VM's Jupyter notebook server, which will allow you to easily execute the code examples while reading along. (Alternatively, you can type the code snippets into the Python interpreter run from an Anaconda prompt.) Instructions for starting the VM's Jupyter server and loading notebook files can be found [here](https://docs.microsoft.com/en-us/azure/machine-learning/machine-learning-data-science-provision-vm#how-to-create-a-strong-password-for-jupyter-and-start-the-notebook-server).\n",
    "\n",
    "Download or clone [this git repository](https://github.com/Azure/Embarrassingly-Parallel-Image-Classification) to the VM. Copy its contents into the `C:\\dsvm\\notebooks` directory, then launch `image_set_preparation.ipynb` from the default starting page of the Jupyter notebook server.\n",
    "\n",
    "If running the code in the Jupyter notebook, set the kernel to \"`conda env:py35`\". If you are not prompted to change the kernel, you can do so by navigating to Kernel -> Change Kernel."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "<a name=\"install\"></a>\n",
    "### Install Python packages and supporting software\n",
    "\n",
    "#### LizardTech's GeoExpress Command Line Applications\n",
    "\n",
    "Download LizardTech's free [GeoExpress Command Line Applications for 64-bit Windows](https://s3.amazonaws.com/bin.us.lizardtech.com/utilities/GeoExpressCLUtils-9.5.0.4326-win64.zip) (available in the [GIS Tools and Utilities](https://www.lizardtech.com/gis-tools/tools-and-utilities) section), which we will use to convert NAIP imagery from the provided [MrSID](https://en.wikipedia.org/wiki/MrSID) format to GeoTIFF format.\n",
    "\n",
    "After downloading and decompressing the archive on your VM, copy the files in the `bin` subdirectory to a new folder on the temporary drive, `D:\\lizardtech`. Define the path to the binary file `mrsidgeodecode` in a Python variable named `mrsidgeodecode_path`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "mrsidgeodecode_path = 'D:\\\\lizardtech\\\\mrsidgeodecode'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "Note that if you are running this notebook on the VM's Jupyter notebook server, you can define this variable by clicking on the code cell and pressing Ctrl+Enter."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "#### Python packages\n",
    "To read the resulting GeoTIFF, we will use the GDAL package from within Python. As of this writing, Christoph Gohlke's [Unofficial Windows Binaries for Python Extension Packages](http://www.lfd.uci.edu/~gohlke/pythonlibs/) provide the easiest means to install the GDAL package, binaries, and library. We recommend that you download the Python 3.5 (`cp35`), `win_amd64` wheels for the following packages from Christoph's site:\n",
    "- [GDAL](http://www.lfd.uci.edu/~gohlke/pythonlibs/#gdal)\n",
    "- [PyPROJ](http://www.lfd.uci.edu/~gohlke/pythonlibs/#pyproj)\n",
    "- [Basemap](http://www.lfd.uci.edu/~gohlke/pythonlibs/#basemap)\n",
    "\n",
    "To install these wheels:\n",
    "1. Launch an Anaconda Prompt.\n",
    "1. Type the following commands (in this order):\n",
    "\n",
    "  ```activate py35\n",
    "  pip install [GDAL wheel filepath]\n",
    "  pip install [PyPROJ wheel filepath]\n",
    "  pip install [Basemap wheel filepath]\n",
    "  ```\n",
    "  \n",
    "  \n",
    "  If you encounter a \"BadZipFile\" error when installing the Basemap wheel, you may need to upgrade `pip` first:\n",
    "  ```\n",
    "  python -m pip install --upgrade pip\n",
    "  ```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "<a name=\"download\"></a>\n",
    "### Download input data locally\n",
    "\n",
    "In this step, you will download the raw imagery and land use labels needed for image extraction. You may skip directly to the [Dataset preparation for deep learning](#prep) section if you prefer to work with the extracted image sets we provide.\n",
    "\n",
    "The input data and intermediate files generated in this tutorial are quite large (>90 GB). We recommend downloading and storing all files in your VM's temporary filespace on the `D` drive: files in this location will be deleted if the machine is restarted, but this will not be problematic if you complete the tutorial in one sitting and store the output images in a more permanent location. Alternatively, you can [provision an additional data disk](https://docs.microsoft.com/en-us/azure/virtual-machines/virtual-machines-windows-attach-disk-portal) to ensure that your files will not be deleted if the VM restarts."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "#### National Land Cover Database\n",
    "\n",
    "[Download](https://www.mrlc.gov/nlcd11_data.php) the compressed 2011 National Land Cover Database file (1.1 GB) and extract its contents. Copy the `nlcd_2011_landcover_2011_edition_2014_10_10` subdirectory to the `D` drive. (Note that multiple nested folders of the same name may be created during extraction; copy the innermost folder.)\n",
    "\n",
    "The main files of interest in this directory are the `.img` image file (erroneously referred to in the Windows Explorer description as a \"Disc Image File\") and its `.ige` \"image extension\". Store the full path to the `.img` file in a Python variable named `nlcd_filepath`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "nlcd_filepath = 'D:\\\\nlcd_2011_landcover_2011_edition_2014_10_10\\\\nlcd_2011_landcover_2011_edition_2014_10_10.img'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "#### National Agriculture Imagery Program\n",
    "\n",
    "NAIP Compressed County Mosaics can be obtained via the [Geospatial Data Gateway](https://gdg.sc.egov.usda.gov/). You may choose any county of interest, but for simplicity, this tutorial will assume that you used Middlesex County, MA 2016 data (which we have [mirrored for your convenience](https://mawahstorage.blob.core.windows.net/aerialimageclassification/naipsample/ortho_imagery_NAIPM16_ma017_3344056_01.zip)).\n",
    "\n",
    "Download and decompress the file, then copy the `ortho_imagery_NAIPM16_ma017_3344056_01` folder to the `D` drive. (Note that multiple nested folders of the same name may have been created; copy the innermost folder.) Store the path to the folder in a Python variable named `naip_dir`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "naip_dir = 'D:\\\\ortho_imagery_NAIPM16_ma017_3344056_01\\\\ortho_imagery'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "The main file of interest in this directory is the `.sid` ([MrSID](https://en.wikipedia.org/wiki/MrSID)) file. This image's data are substantially compressed in this format: they will grow in size by ~40x when we convert the image to GeoTIFF format for further analysis. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "<a name=\"production\"></a>\n",
    "## Produce image sets for training and evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "<a name=\"overview\"></a>\n",
    "### Overview of the approach\n",
    "\n",
    "Our goal is to extract many smaller aerial images (\"tiles\") from the NAIP image of the entire county. Each tile will be labeled with a land use classification determined from the NLCD data. Before launching into the technical details of this process, we provide an overview of the steps involved."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "#### Choosing tile size\n",
    "ResNet DNNs are traditionally trained with 224 px x 224 px images: this sets our minimum tile size at 224 meters x 224 meters (because our aerial imagery has approx. 1 meter resolution). A sample tile with these dimensions is shown below at full size.\n",
    "\n",
    "<img src=\"./img/extraction/sample_tile.png\" />"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "#### Proposing candidate tiles\n",
    "\n",
    "We divided the Middlesex County, MA region along a rectangular grid of tiles, as illustrated in the approximately 1 km x 1 km region shown below: \n",
    "\n",
    "<img src=\"./img/extraction/common_naip_tiled.png\" />\n",
    "\n",
    "To create the grid, we will need to find the latitude/longitude boundaries and dimensions in meters of the county-level NAIP image. We do this by converting the image to GeoTIFF format and using the Python GDAL package to extract the necessary information from the header. Once we have defined the grid, we can use GDAL to extract the image for each tile from the county-level image."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "#### Assigning land use classes to tiles\n",
    "\n",
    "The tiles described above often contain multiple types of land, making it difficult to assign a single land use label. To illustrate this problem, we overlaid the NLCD land use labels and the NAIP aerial image in a region centered on the [Boston Common](https://binged.it/2kTjwOe) (NAIP-only image shown above):\n",
    "\n",
    "<img src=\"./img/extraction/common_tiled_only.png\" />"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "A straightforward approach to avoid ambiguity when assigning labels would be to remove any tiles with multiple land uses. Unfortunately, we found this requirement too restrictive: the number of tiles retained by this method was too small for model training. We developed an arbitrary, more lenient criterion for labeling in which we sampled the land use labels at nine regularly-spaced points within each tile, as depicted below:\n",
    "\n",
    "<img src=\"./img/extraction/common_points.png\" />\n",
    "\n",
    "If all nine points had the same land use class, we assigned that land use class as the tile's label. Images that could not be assigned a label by this method were not included in the training and validation sets."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "<a name=\"conversion\"></a>\n",
    "### Convert NAIP images from MrSID to GeoTIFF\n",
    "\n",
    "We used [LizardTech's GeoExpress Command Line Application](LizardTech's GeoExpress Command Line Application) `mrsidgeodecode` to convert the aerial imagery data from MrSID format to GeoTIFF. This command may take several minutes to run. The resulting GeoTIFF file is roughly 45 GB in size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import subprocess\n",
    "\n",
    "filename_base = None\n",
    "for filename in os.listdir(naip_dir):\n",
    "    if filename.endswith('.sid'):\n",
    "        filename_base = filename.split('.sid')[0]\n",
    "        break\n",
    "if filename_base == None:\n",
    "    raise Exception(\"Couldn't find the MrSID file in folder {}; is the name correct?\".format(naip_dir))\n",
    "    \n",
    "geodecode_command = [mrsidgeodecode_path, '-wf',\n",
    "                     '-i', os.path.join(naip_dir, '{}.sid'.format(filename_base)),\n",
    "                     '-o', os.path.join(naip_dir, '{}.tif'.format(filename_base))]\n",
    "subprocess.check_output(geodecode_command)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "If you encounter a CalledProcessError when running the code cell above, you may need to update the write permissions for the `D` drive:\n",
    "- Right-click on the \"Temporary storage (`D:`)\" drive under \"This PC\" in Windows Explorer and select \"Properties'.\n",
    "- Under the Security tab, click the 'Edit\" button.\n",
    "- Highlight the group name \"Everyone\" and click the \"Allow\" checkbox next to \"Full control\". Other checkboxes will be automatically populated.\n",
    "- Click the \"Apply\" button. You may receive an error that `pagefile.sys` cannot be updated; this can be disregarded.\n",
    "- When permission updates are complete, click OK.\n",
    "- Reattempt to execute the code above."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "<a name=\"tiling\"></a>\n",
    "### Add tiling for fast image extraction\n",
    "\n",
    "To speed up extraction of images from arbitrary regions in the 45 GB file, we add tiling to the image. Since we must write the tiled image before deleting the untiled version, we briefly double our disk space usage at this step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "gdal_command = ['C:\\\\Anaconda\\\\envs\\\\py35\\\\Lib\\\\site-packages\\\\osgeo\\\\gdal_translate',\n",
    "            '-co', 'TILED=YES',\n",
    "            os.path.join(naip_dir, '{}.tif'.format(filename_base)),\n",
    "            os.path.join(naip_dir, '{}_tiled.tif'.format(filename_base))]\n",
    "subprocess.check_output(gdal_command)\n",
    "naip_filepath = os.path.join(naip_dir, '{}_tiled.tif'.format(filename_base))\n",
    "\n",
    "# Remove the large intermediate file if the command ran successfully\n",
    "if os.path.exists(naip_filepath):\n",
    "    os.remove(os.path.join(naip_dir, '{}.tif'.format(filename_base)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "<a name=\"box\"></a>\n",
    "### Find a lat-lon bounding box for the region\n",
    "\n",
    "We now extract the county's boundaries in latitude/longitude coordinates from the GeoTIFF file. Note that because of the county's irregular shape, some regions inside this lat/lon bounding box will not contain any data. The bounding box is nonetheless useful for establishing the dimensions of a tiling grid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "from osgeo import gdal\n",
    "from gdalconst import *\n",
    "import osr\n",
    "from mpl_toolkits.basemap import Basemap\n",
    "from collections import namedtuple\n",
    "\n",
    "LatLonBounds = namedtuple('LatLonBounds', ['llcrnrlat', 'llcrnrlon', 'urcrnrlat', 'urcrnrlon'])\n",
    "\n",
    "def get_bounding_box(naip_filepath):\n",
    "    ''' Finds a bounding box for the NAIP GeoTIFF in lat/lon '''\n",
    "    naip_image = gdal.Open(naip_filepath, GA_ReadOnly)\n",
    "    naip_proj = osr.SpatialReference()\n",
    "    naip_proj.ImportFromWkt(naip_image.GetProjection())\n",
    "    naip_ulcrnrx, naip_xstep, _, naip_ulcrnry, _, naip_ystep = naip_image.GetGeoTransform()\n",
    "\n",
    "    world_map = Basemap(lat_0=0,\n",
    "                        lon_0=0,\n",
    "                        llcrnrlat=-90, urcrnrlat=90,\n",
    "                        llcrnrlon=-180, urcrnrlon=180,\n",
    "                        resolution='c', projection='stere')\n",
    "    world_proj = osr.SpatialReference()\n",
    "    world_proj.ImportFromProj4(world_map.proj4string)\n",
    "    ct_to_world = osr.CoordinateTransformation(naip_proj, world_proj)\n",
    "    \n",
    "    lats = []\n",
    "    lons = []\n",
    "    for corner_x, corner_y in [(naip_ulcrnrx, naip_ulcrnry),\n",
    "                               (naip_ulcrnrx, naip_ulcrnry + naip_image.RasterYSize * naip_ystep),\n",
    "                               (naip_ulcrnrx + naip_image.RasterXSize * naip_xstep,\n",
    "                                naip_ulcrnry + naip_image.RasterYSize * naip_ystep),\n",
    "                               (naip_ulcrnrx + naip_image.RasterXSize * naip_xstep, naip_ulcrnry)]:\n",
    "        xpos, ypos, _ = ct_to_world.TransformPoint(corner_x, corner_y)\n",
    "        lon, lat = world_map(xpos, ypos, inverse=True)\n",
    "        lats.append(lat)\n",
    "        lons.append(lon)\n",
    "\n",
    "    return(LatLonBounds(llcrnrlat=min(lats),\n",
    "                        llcrnrlon=min(lons),\n",
    "                        urcrnrlat=max(lats),\n",
    "                        urcrnrlon=max(lons)))\n",
    "\n",
    "region_bounds = get_bounding_box(naip_filepath)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "<a name=\"meters\"></a>\n",
    "### Get the approximate dimensions of the bounding box in meters\n",
    "\n",
    "Here we approximate the width/height of the bounding box in meters, using the fact that latitude does not change substantially on the county scale. This dimensions in meters are needed later to define a grid of 224 meter x 224 meter tiles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "RegionSize = namedtuple('RegionSize', ['width', 'height'])  # in meters!\n",
    "import numpy as np\n",
    "\n",
    "def get_approx_region_size(region_bounds):\n",
    "    ''' Returns the region width (at mid-lat) and height in meters'''\n",
    "    mid_lat_radians = (region_bounds.llcrnrlat + region_bounds.urcrnrlat) * \\\n",
    "                      (np.pi / 360)\n",
    "    earth_circumference = 6.371E6 * 2 * np.pi # in meters\n",
    "    region_middle_width_meters = (region_bounds.urcrnrlon - region_bounds.llcrnrlon) * \\\n",
    "                                 earth_circumference * np.cos(mid_lat_radians) / 360\n",
    "    region_height_meters = (region_bounds.urcrnrlat - region_bounds.llcrnrlat) * \\\n",
    "                           earth_circumference / 360\n",
    "    return(RegionSize(region_middle_width_meters, region_height_meters))\n",
    "\n",
    "approx_region_size = get_approx_region_size(region_bounds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "<a name=\"consistent\"></a>\n",
    "### Find tiles with consistent land use class\n",
    "\n",
    "#### Define helper functions\n",
    "\n",
    "The function below defines helper functions that help us access data from both input images (the NAIP aerial image and the NLCD land use label image) using latitude and longitude coordinate systems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "\n",
    "def create_helper_functions(region_bounds, nlcd_filepath, naip_filepath):\n",
    "    ''' Makes helper functions to label points (NLCD) and extract tiles (NAIP) '''\n",
    "    nlcd_image = gdal.Open(nlcd_filepath, GA_ReadOnly)\n",
    "    nlcd_proj = osr.SpatialReference()\n",
    "    nlcd_proj.ImportFromWkt(nlcd_image.GetProjection())\n",
    "    nlcd_ulcrnrx, nlcd_xstep, _, nlcd_ulcrnry, _, nlcd_ystep = nlcd_image.GetGeoTransform()\n",
    "    \n",
    "    naip_image = gdal.Open(naip_filepath, GA_ReadOnly)\n",
    "    naip_proj = osr.SpatialReference()\n",
    "    naip_proj.ImportFromWkt(naip_image.GetProjection())\n",
    "    naip_ulcrnrx, naip_xstep, _, naip_ulcrnry, _, naip_ystep = naip_image.GetGeoTransform()\n",
    "    \n",
    "    region_map = Basemap(lat_0 = (region_bounds.llcrnrlat + region_bounds.urcrnrlat)/2,\n",
    "                         lon_0 = (region_bounds.llcrnrlon + region_bounds.urcrnrlon)/2,\n",
    "                         llcrnrlat=region_bounds.llcrnrlat,\n",
    "                         llcrnrlon=region_bounds.llcrnrlon,\n",
    "                         urcrnrlat=region_bounds.urcrnrlat,\n",
    "                         urcrnrlon=region_bounds.urcrnrlon,\n",
    "                         resolution='c',\n",
    "                         projection='stere')\n",
    "    \n",
    "    region_proj = osr.SpatialReference()\n",
    "    region_proj.ImportFromProj4(region_map.proj4string)\n",
    "    ct_to_nlcd = osr.CoordinateTransformation(region_proj, nlcd_proj)\n",
    "    ct_to_naip = osr.CoordinateTransformation(region_proj, naip_proj)\n",
    "\n",
    "    def get_nlcd_label(point):\n",
    "        ''' Project lat/lon point to NLCD GeoTIFF; return label of that point '''\n",
    "        basemap_coords = region_map(point.lon, point.lat)  # NB unusual argument order\n",
    "        x, y, _ = [int(i) for i in ct_to_nlcd.TransformPoint(*basemap_coords)]\n",
    "        xoff = int(round((x - nlcd_ulcrnrx) / nlcd_xstep))\n",
    "        yoff = int(round((y - nlcd_ulcrnry) / nlcd_ystep))\n",
    "        label = int(nlcd_image.ReadAsArray(xoff=xoff, yoff=yoff, xsize=1, ysize=1))\n",
    "        return(label)\n",
    "        \n",
    "    def get_naip_tile(tile_bounds, tile_size):\n",
    "        ''' Check that tile lies within county bounds; if so, extract its image '''\n",
    "        \n",
    "        # Transform tile bounds in lat/lon to NAIP projection coordinates\n",
    "        xmax, ymax = region_map(tile_bounds.urcrnrlon, tile_bounds.urcrnrlat)\n",
    "        xmin, ymin = region_map(tile_bounds.llcrnrlon, tile_bounds.llcrnrlat)\n",
    "        xstep = (xmax - xmin) / tile_size.width\n",
    "        ystep = (ymax - ymin) / tile_size.height\n",
    "\n",
    "        grid = np.mgrid[xmin:xmax:tile_size.width * 1j, ymin:ymax:tile_size.height * 1j]\n",
    "        shape = grid[0, :, :].shape\n",
    "        size = grid[0, :, :].size\n",
    "        xy_target = np.array(ct_to_naip.TransformPoints(grid.reshape(2, size).T))\n",
    "        xx = xy_target[:,0].reshape(shape)\n",
    "        yy = xy_target[:,1].reshape(shape)\n",
    "        \n",
    "        # Extract rectangle from NAIP GeoTIFF containing superset of needed points\n",
    "        xoff = int(round((xx.min() - naip_ulcrnrx) / naip_xstep))\n",
    "        yoff = int(round((yy.max() - naip_ulcrnry) / naip_ystep))\n",
    "        xsize_to_use = int(np.ceil((xx.max() - xx.min())/np.abs(naip_xstep))) + 1\n",
    "        ysize_to_use = int(np.ceil((yy.max() - yy.min())/np.abs(naip_ystep))) + 1\n",
    "        data = naip_image.ReadAsArray(xoff=xoff,\n",
    "                                      yoff=yoff,\n",
    "                                      xsize=xsize_to_use,\n",
    "                                      ysize=ysize_to_use)        \n",
    "        # Map the pixels of interest in NAIP GeoTIFF to the tile (might involve rotation or scaling)\n",
    "        image = np.zeros((xx.shape[1], xx.shape[0], 3)).astype(int)  # rows are height, cols are width, third dim is color\n",
    "        \n",
    "        try:\n",
    "            for i in range(xx.shape[0]):\n",
    "                for j in range(xx.shape[1]):\n",
    "                    x_idx = int(round((xx[i,j] - naip_ulcrnrx) / naip_xstep)) - xoff\n",
    "                    y_idx = int(round((yy[i,j] - naip_ulcrnry) / naip_ystep)) - yoff\n",
    "                    image[xx.shape[1] - j - 1, i, :] = data[:, y_idx, x_idx]\n",
    "        except TypeError as e:\n",
    "            # The following can occur if our pixel superset request exceeds the GeoTIFF's bounds\n",
    "            return(None)\n",
    "        \n",
    "        if np.sum(image.sum(axis=2) == 0) > 10: # too many nodata pixels\n",
    "            return None\n",
    "        \n",
    "        image = Image.fromarray(image.astype('uint8'))\n",
    "        return(image)\n",
    "    \n",
    "    return(get_nlcd_label, get_naip_tile)\n",
    "get_nlcd_label, get_naip_tile = create_helper_functions(region_bounds, nlcd_filepath, naip_filepath)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "#### Tile selection\n",
    "\n",
    "We form a rectangular grid of candidate tiles spanning as much of the bounding box as possible. The boundaries of tiles that could be assigned a label are noted, and will later be used to extract the corresponding aerial image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 46999 tiles to extract\n"
     ]
    }
   ],
   "source": [
    "LatLonPosition = namedtuple('LatLonPosition', ['lat', 'lon'])\n",
    "Tile = namedtuple('Tile', ['bounds', 'label'])\n",
    "\n",
    "nlcd_label_to_class = {21: 'Developed',\n",
    "                       22: 'Developed',\n",
    "                       23: 'Developed',\n",
    "                       24: 'Developed',\n",
    "                       31: 'Barren',\n",
    "                       41: 'Forest',\n",
    "                       42: 'Forest',\n",
    "                       43: 'Forest',\n",
    "                       90: 'Forest',\n",
    "                       51: 'Shrub',\n",
    "                       52: 'Shrub',\n",
    "                       71: 'Herbaceous',\n",
    "                       72: 'Herbaceous',\n",
    "                       73: 'Herbaceous',\n",
    "                       74: 'Herbaceous',\n",
    "                       95: 'Herbaceous',\n",
    "                       81: 'Cultivated',\n",
    "                       82: 'Cultivated'}\n",
    "\n",
    "def find_tiles_with_consistent_labels(region_bounds, region_size, tile_size):\n",
    "    ''' Find tiles for which nine grid points all have the same label '''\n",
    "    tiles_wide = int(np.floor(region_size.width / tile_size.width))\n",
    "    tiles_tall = int(np.floor(region_size.height / tile_size.height))\n",
    "    tile_width = (region_bounds.urcrnrlon - region_bounds.llcrnrlon) / (region_size.width / tile_size.width)\n",
    "    tile_height = (region_bounds.urcrnrlat - region_bounds.llcrnrlat) / (region_size.height / tile_size.height)\n",
    "    \n",
    "    current_lat = region_bounds.llcrnrlat\n",
    "    current_lon = region_bounds.llcrnrlon\n",
    "    \n",
    "    tiles_to_use = []\n",
    "    for i in range(tiles_tall):\n",
    "        for j in range(tiles_wide):\n",
    "            try:\n",
    "                labels = []\n",
    "                for k in range(3):\n",
    "                    for ell in range(3):\n",
    "                        my_label = get_nlcd_label(LatLonPosition(lat=current_lat + tile_width * (1 + 2*k) / 6,\n",
    "                                                                 lon=current_lon + tile_height * (1 + 2*ell) / 6))\n",
    "                        labels.append(nlcd_label_to_class[my_label])\n",
    "                num_matching = np.sum(np.array(labels) == labels[4])\n",
    "                if (num_matching == 9):\n",
    "                    bounds = LatLonBounds(llcrnrlat=current_lat,\n",
    "                                          llcrnrlon=current_lon,\n",
    "                                          urcrnrlat=current_lat + tile_height,\n",
    "                                          urcrnrlon=current_lon + tile_width)\n",
    "                    tiles_to_use.append(Tile(bounds=bounds,\n",
    "                                             label=labels[4]))\n",
    "            except KeyError:\n",
    "                pass\n",
    "            current_lon += tile_width\n",
    "        current_lon = region_bounds.llcrnrlon\n",
    "        current_lat += tile_height\n",
    "    return(tiles_to_use)\n",
    "\n",
    "tile_size = RegionSize(224, 224)\n",
    "tiles = find_tiles_with_consistent_labels(region_bounds, approx_region_size, tile_size)\n",
    "print('Found {} tiles to extract'.format(len(tiles)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "Note that we have grouped together similar NLCD land cover labels:\n",
    "- **Developed**: 21, 22, 23, 24\n",
    "- **Forest**: 41, 42, 43, 90\n",
    "- **Barren**: 31\n",
    "- **Shrub**: 51, 52\n",
    "- **Herbaceous**: 71, 72, 73, 74, 95\n",
    "- **Cultivated**: 81, 82"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "<a name=\"extraction\"></a>\n",
    "### Extract and save images for tiles of interest\n",
    "\n",
    "The most time-consuming step is the extraction of images from the GeoTIFF. We use a helper function defined earlier, `naip_get_tile`, to check that the tile lies entirely inside the county shapefile and, if so, return the extracted image. Images are sorted into directories based on land use class. Expect this step to take several hours. You can monitor the progress by checking the number of images created in the output directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "def extract_tiles(tiles, dest_folder, filename_base):\n",
    "    ''' Coordinates saving tile data, including extracted images and CSV descriptions '''\n",
    "    if not os.path.exists(dest_folder):\n",
    "        os.makedirs(dest_folder)\n",
    "        \n",
    "    tile_descriptions = []\n",
    "    i = 0\n",
    "    while i < len(tiles):\n",
    "        tile = tiles[i]\n",
    "        tile_image = get_naip_tile(tile.bounds, tile_size)\n",
    "        i += 1\n",
    "        if (tile_image is None):\n",
    "            continue  # tile did not lie entirely within the county boundary (it was at least partially blank)\n",
    "        my_directory = os.path.join(dest_folder, '{}'.format(tile.label))\n",
    "        my_filename = os.path.join(my_directory, '{}_{}.png'.format(filename_base, i))\n",
    "        if not os.path.exists(my_directory):\n",
    "            os.makedirs(my_directory)\n",
    "        tile_image.save(my_filename, 'PNG')\n",
    "    return\n",
    "\n",
    "output_dir = 'D:\\\\tiles'\n",
    "extract_tiles(tiles, output_dir, filename_base)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "<a name=\"partition\"></a>\n",
    "### Dataset partitioning\n",
    "\n",
    "To create large and diverse training/validation sets, we repeated the steps above for 12 counties spread across the United States. We partitioned the images into training and validation sets at the county level: this division allows a reasonably accurate estimate of model performance by ensuring some dissimilarity between images seen during training and those used to evaluate the model's performance. We balanced each image set by randomly removing images from overrepresented classes.\n",
    "\n",
    "While you are welcome to repeat those steps yourself, for expediency we recommend that you continue with this tutorial by downloading our training and validation sets (linked in the next section)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "<a name=\"prep\"></a>\n",
    "## Prepare deep learning framework-specific input files\n",
    "\n",
    "If you have not generated your own training and validation sets through image extraction, download the following files and decompress them in your VM's `D` drive:\n",
    "- [Balanced training image set (~3 GB)](https://mawahstorage.blob.core.windows.net/aerialimageclassification/imagesets/balanced_training_set.zip)\n",
    "- [Balanced validation image set (~1 GB)](https://mawahstorage.blob.core.windows.net/aerialimageclassification/imagesets/balanced_validation_set.zip)\n",
    "\n",
    "The image sets linked above contain raw PNG images sorted into folders by their assigned label. Many deep learning frameworks require proprietary image formats or supporting files to efficiently load images in minibatches for training. We now produce the supporting files needed by our CNTK and TensorFlow training scripts. Note that we do not need to produce similar supporting files for the validation image set, because we will not use minibatching when applying the trained models to the validation set.\n",
    "\n",
    "Update the `training_image_dir` variable below to reflect the directory where your training and validation sets have been saved. The `label_to_number_dict` variable specifies the correspondence between the label names and a numeric code; it does not need to be modified unless you have changed the labeling scheme."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "\n",
    "training_image_dir = 'D:\\\\balanced_training_set'\n",
    "label_to_number_dict = {'Barren': 0,\n",
    "                        'Forest': 1,\n",
    "                        'Shrub': 2,\n",
    "                        'Cultivated': 3,\n",
    "                        'Herbaceous': 4,\n",
    "                        'Developed': 5}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "<a name=\"cntk\"></a>\n",
    "### Cognitive Toolkit (CNTK)\n",
    "\n",
    "Our CNTK training script uses a MAP file -- a tab-delimited file where each line lists an image's filepath and label -- to load image data in minibatches. We generate the MAP file as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(training_image_dir, 'map.txt'), 'w') as map_file:\n",
    "    for label in np.sort(os.listdir(training_image_dir)):\n",
    "        my_dir = os.path.join(training_image_dir, label)\n",
    "        if not os.path.isdir(my_dir):\n",
    "            continue\n",
    "        for filename in os.listdir(my_dir):\n",
    "            map_file.write('{}\\t{}\\n'.format(os.path.join(my_dir, filename),\n",
    "                                             label_to_number_dict[label]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "<a name=\"tf\"></a>\n",
    "### TensorFlow\n",
    "\n",
    "We made use of the [`tf-slim` API](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/slim) for TensorFlow, which provides pre-trained ResNet models and helpful scripts for retraining and scoring. These scripts require converting the image set into [TFRecords](https://www.tensorflow.org/how_tos/reading_data/#file_formats) for minibatching. (Each TFRecord contains many image files as well as their labels.) We also create a `labels.txt` file mapping the labels to integer values, and a `dataset_split_info.csv` file describing the images assigned to the training set.\n",
    "\n",
    "The following code was modified from the [TensorFlow models repo's slim subdirectory](https://github.com/tensorflow/models/tree/master/slim)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 44184 training images\n"
     ]
    }
   ],
   "source": [
    "# Original Copyright 2016 The TensorFlow Authors. All Rights Reserved.\n",
    "# Modified 2017 by Microsoft Corporation.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "# ==============================================================================\n",
    "np.random.seed(5318)\n",
    "\n",
    "class ImageReader(object):\n",
    "    def __init__(self):\n",
    "        # Initializes function that decodes RGB JPEG data.\n",
    "        self._decode_png_data = tf.placeholder(dtype=tf.string)\n",
    "        self._decode_png = tf.image.decode_png(self._decode_png_data, channels=3)\n",
    "\n",
    "    def read_image_dims(self, sess, image_data):\n",
    "        image = self.decode_png(sess, image_data)\n",
    "        return image.shape[0], image.shape[1]\n",
    "\n",
    "    def decode_png(self, sess, image_data):\n",
    "        image = sess.run(self._decode_png,\n",
    "                         feed_dict={self._decode_png_data: image_data})\n",
    "        assert len(image.shape) == 3\n",
    "        assert image.shape[2] == 3\n",
    "        return image\n",
    "    \n",
    "def image_to_tfexample(image_data, image_format, height, width, class_id):\n",
    "    return tf.train.Example(features=tf.train.Features(feature={\n",
    "        'image/encoded': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image_data])),\n",
    "        'image/format': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image_format])),\n",
    "        'image/class/label': tf.train.Feature(int64_list=tf.train.Int64List(value=[class_id])),\n",
    "        'image/height': tf.train.Feature(int64_list=tf.train.Int64List(value=[height])),\n",
    "        'image/width': tf.train.Feature(int64_list=tf.train.Int64List(value=[width])),\n",
    "    }))\n",
    "\n",
    "def find_images(image_dir):\n",
    "    training_filenames = []\n",
    "    for folder in os.listdir(image_dir):\n",
    "        folder_path = os.path.join(image_dir, folder)\n",
    "        if not os.path.isdir(folder_path):\n",
    "            continue\n",
    "        ''' This is a new directory/label -- consider all images inside it '''\n",
    "        my_filenames = []\n",
    "        for filename in os.listdir(folder_path):\n",
    "            my_filenames.append(os.path.join(folder_path, filename))\n",
    "        training_filenames.extend(my_filenames)\n",
    "    print('Found {} training images'.format(len(training_filenames)))\n",
    "    return(training_filenames)\n",
    "    \n",
    "def write_dataset(dataset_name, split_name, my_filenames,  image_dir, n_shards=5):\n",
    "    num_per_shard = int(np.ceil(len(my_filenames) / n_shards))\n",
    "    records = []\n",
    "    with tf.Graph().as_default():\n",
    "        image_reader = ImageReader()\n",
    "        with tf.Session('') as sess:\n",
    "            for shard_idx in range(n_shards):\n",
    "                shard_filename = os.path.join(image_dir,\n",
    "                                              '{}_{}_{:05d}-of-{:05d}.tfrecord'.format(dataset_name,\n",
    "                                                                                       split_name,\n",
    "                                                                                       shard_idx+1,\n",
    "                                                                                       n_shards))\n",
    "                with tf.python_io.TFRecordWriter(shard_filename) as tfrecord_writer:\n",
    "                    for image_idx in range(num_per_shard * shard_idx,\n",
    "                                           min(num_per_shard * (shard_idx+1), len(my_filenames))):\n",
    "                        with open(my_filenames[image_idx], 'rb') as f:\n",
    "                            image_data = f.read()\n",
    "                        height, width = image_reader.read_image_dims(sess, image_data)\n",
    "                        class_name = os.path.basename(os.path.dirname(my_filenames[image_idx]))\n",
    "                        class_id = label_to_number_dict[class_name]\n",
    "                        example = image_to_tfexample(image_data, b'png', height, width, class_id)\n",
    "                        tfrecord_writer.write(example.SerializeToString())\n",
    "                        records.append([dataset_name, split_name, my_filenames[image_idx], shard_idx,\n",
    "                                        image_idx, class_name, class_id])\n",
    "    df = pd.DataFrame(records, columns=['dataset_name', 'split_name', 'filename', 'shard_idx', 'image_idx',\n",
    "                                        'class_name', 'class_id'])\n",
    "    return(df)\n",
    " \n",
    "training_image_dir = 'D:\\\\balanced_training_set'\n",
    "training_filenames = find_images(training_image_dir)\n",
    "training_filenames = np.random.permutation(training_filenames)\n",
    "df = write_dataset('aerial', 'train', training_filenames, training_image_dir, n_shards=50)\n",
    "df.to_csv(os.path.join(training_image_dir, 'dataset_split_info.csv'), index=False)\n",
    "\n",
    "with open(os.path.join(training_image_dir, 'labels.txt'), 'w') as f:\n",
    "    for key, value in label_to_number_dict.items():\n",
    "        f.write('{}:{}\\n'.format(key, value))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "## Next Steps\n",
    "\n",
    "* To retrain image classification DNNs using the training image set, proceed to the [Model Training](https://github.com/Azure/Embarrassingly-Parallel-Image-Classification/blob/master/model_training.ipynb) notebook.\n",
    "   * You may skip this step if you choose to use our example retrained DNNs.\n",
    "* To apply DNNs to the validation set images on Spark, proceed to the [Scoring on Spark](./scoring_on_spark.ipynb) notebook."
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}