notebooks/xview/xview-vehicles-data-prep.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# xView Vehicle Object Detection Data Prep\n",
    "\n",
    "This notebook prepares data for training an object detection model on the xView dataset.\n",
    "\n",
    "\n",
    "* Download the training images and labels from the xView competition site, unzip them, and put the contents of each zipfile in a local or S3 directory.\n",
    "* Set `raw_uri` to this directory containing the raw dataset.\n",
    "* Set `processed_uri` to a local or S3 directory (you can write to), which will store the processed data generated by this notebook.\n",
    "\n",
    "This is all you will need to do in order to run this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# raw_uri = '/opt/data/examples/xview/raw-data'\n",
    "raw_uri = 's3://raster-vision-xview-example/raw-data'\n",
    "processed_uri = '/opt/data/examples/xview/processed-data'\n",
    "# processed_uri = 's3://raster-vision-xview-example/processed-data'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The steps we'll take to prepare the data are as follows:\n",
    "\n",
    "- Filter out all of the non-vehicle bounding boxes from the labels. Combine all vehicle types into one class. \n",
    "- Subset the entire xView dataset to only include the images that are most densely populated with vehicles.\n",
    "- Split the selected images randomly into 80%/20% training and validation sets\n",
    "- Split the vehicle labels by image, and save off a label GeoJSON file per image\n",
    "\n",
    "\n",
    "This process will save the split labels, and `train_scenes.csv` and `val_scenes.csv` files that are used by the experiment at `examples.xview.object_detection` to `processed_uri`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from os.path import join\n",
    "import json\n",
    "import random\n",
    "from collections import defaultdict\n",
    "\n",
    "from rastervision.utils.files import (\n",
    "    download_if_needed, list_paths, file_to_json, json_to_file, \n",
    "    get_local_path, make_dir, sync_to_dir, str_to_file)\n",
    "\n",
    "random.seed(12345)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Filter out non-vehicle labels\n",
    "\n",
    "The xView dataset includes labels for a number of different types of objects. We are only interested in building a detector for objects that can be categorized as vehicles (e.g. 'small car', 'passenger vehicle', 'bus'). We have pre-determined the ids that map to vehicle labels and will use them to extract all the vehicles from the whole xView label set. In this section we also assign a class name of 'vehicle' to all of the resulting labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "label_uri = join(raw_uri, 'xView_train.geojson')\n",
    "label_js = file_to_json(label_uri)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "vehicle_type_ids = [17, 18, 19, 20, 21, 23, 24, 25, 26, 27, 28, 29, 32, \n",
    "                    53, 54, 55, 56, 57, 59, 60, 61, 62, 63, 64, 65, 66]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "vehicle_features = []\n",
    "for f in label_js['features']:\n",
    "    if f['properties']['type_id'] in vehicle_type_ids:\n",
    "        f['properties']['class_name'] = 'vehicle'\n",
    "        vehicle_features.append(f)\n",
    "label_js['features'] = vehicle_features"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Subset images with the most vehicles\n",
    "\n",
    "In this section we determine which images contain the most vehicles and are therefore the best candidates for this experiment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_to_vehicle_counts = defaultdict(int)\n",
    "for f in label_js['features']:\n",
    "    image_id = f['properties']['image_id']\n",
    "    image_to_vehicle_counts[image_id] += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use top 10% of images by vehicle count.\n",
    "experiment_image_count = round(len(image_to_vehicle_counts.keys()) * 0.1)\n",
    "sorted_images_and_counts = sorted(image_to_vehicle_counts.items(), key=lambda x: x[1])\n",
    "selected_images_and_counts = sorted_images_and_counts[-experiment_image_count:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Split into train and validation\n",
    "\n",
    "Split up training and validation data. Use 80% of images in the training set and 20% in the validation set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "ratio = 0.8\n",
    "training_sample_size = round(ratio * experiment_image_count)\n",
    "train_sample = random.sample(range(experiment_image_count), training_sample_size)\n",
    "\n",
    "train_images = []\n",
    "val_images = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(training_sample_size):\n",
    "    img = selected_images_and_counts[i][0]\n",
    "    img_path = join('train_images', img)\n",
    "    if i in train_sample:\n",
    "        train_images.append(img_path)\n",
    "    else:\n",
    "        val_images.append(img_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Divide labels up by image\n",
    "\n",
    "Using one vehicle label geojson for all of the training and validation images can become unwieldy. Instead, we will divide the labels up so that each image has a unique geojson associated with it. We will save each of these geojsons to the base directory you provided at the outset.\n",
    "\n",
    "Then, we will create CSVs that our experiments will use to load the training and validation data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def subset_labels(images):\n",
    "    for i in images:\n",
    "        img_fn = os.path.basename(i)\n",
    "        img_id = os.path.splitext(img_fn)[0]\n",
    "        tiff_features = []\n",
    "        for l in label_js['features']:\n",
    "            image_id = l['properties']['image_id']\n",
    "            if image_id == img_fn:\n",
    "                tiff_features.append(l)\n",
    "\n",
    "        tiff_geojson = {}\n",
    "        for key in label_js:\n",
    "            if not key == 'features':\n",
    "                tiff_geojson[key] = label_js[key]\n",
    "        tiff_geojson['features'] = tiff_features\n",
    "        \n",
    "        json_to_file(tiff_geojson, join(processed_uri, 'labels', '{}.geojson'.format(img_id)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "subset_labels(train_images)\n",
    "subset_labels(val_images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_csv(images, path):\n",
    "    csv_rows = []\n",
    "    for img in images:\n",
    "        img_id = os.path.splitext(os.path.basename(img))[0]\n",
    "        img_path = join('train_images', '{}.tif'.format(img_id))\n",
    "        labels_path = join('labels','{}.geojson'.format(img_id))\n",
    "        csv_rows.append('\"{}\",\"{}\"'.format(img_path, labels_path))\n",
    "    str_to_file('\\n'.join(csv_rows), path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "create_csv(train_images, join(processed_uri, 'train-scenes.csv'))\n",
    "create_csv(val_images, join(processed_uri, 'val-scenes.csv'))"
   ]
  },
  {
   "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}