segmentation-model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e78e059b",
   "metadata": {},
   "source": [
    "## Train a Segmentation Model\n",
    "\n",
    "This notebook trains a segmentation model (using [segmentation-models-pytorch](https://smp.readthedocs.io/en/latest/index.html)) on some [Sentinel-2](https://planetarycomputer.microsoft.com/dataset/sentinel-2-l2a) data using crop labels from Radiant Earth's [South Africa Crop Type Competition](http://registry.mlhub.earth/10.34911/rdnt.j0co8q/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7549d67d-52bf-4d0e-acc4-84d7786c09a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pystac\n",
    "import pystac_client\n",
    "import requests\n",
    "import shapely.geometry\n",
    "import rioxarray\n",
    "import planetary_computer\n",
    "import stackstac\n",
    "import torch\n",
    "import torch.utils.data\n",
    "import rasterio.plot\n",
    "import numpy as np\n",
    "import dask\n",
    "\n",
    "import contextlib\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\", \"Creating an ndarray from ragged\")\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    device = \"cuda:0\"\n",
    "    device_ctx = torch.cuda.device\n",
    "else:\n",
    "    device = \"cpu\"\n",
    "    device_ctx = contextlib.nullcontext\n",
    "    # avoid some warnings from the pynvml monitor\n",
    "    import subprocess\n",
    "\n",
    "    subprocess.call([\"pip\", \"uninstall\", \"-y\", \"pynvml\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fab048a-b85d-487f-a4a3-dddd17d23cad",
   "metadata": {},
   "source": [
    "We'll use the same STAC catalog for the labels that we used previously."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f6c787a-e7b5-4654-acb1-1ae8d037cc09",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_catalog = pystac.read_file(\n",
    "    \"https://esip2021.blob.core.windows.net/esip2021/train/collection.json\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d15e36f",
   "metadata": {},
   "source": [
    "And we'll process the data in parallel (on a single machine) using [Dask](https://dask.org/). We won't use too many features of Dask, we'll just use [dask.delayed](https://docs.dask.org/en/latest/delayed.html) to run some functions in parallel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb3dc988-b22a-4ebc-bab7-5e296266c55b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client\n",
    "\n",
    "client = Client()\n",
    "client.run(lambda: warnings.filterwarnings(\"ignore\", \"Creating an ndarray from ragged\"))\n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "291ad1f7",
   "metadata": {},
   "source": [
    "Make sure to open that link in the dashboard or dask labextension."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25eb0518",
   "metadata": {},
   "source": [
    "## Load Training STAC Items\n",
    "\n",
    "In the last notebook, we worked with a single STAC item / scene for the `labels`. To better generalize to more fields, we'll now load a bunch of them. To cut down on the time we spend waiting for data loading and training, we'll work with just the first 100 items."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c99e9b9-45d8-477f-948e-165e6816bab2",
   "metadata": {},
   "outputs": [],
   "source": [
    "N_SCENES = 100\n",
    "links = training_catalog.get_item_links()[:N_SCENES]\n",
    "label_items = [link.resolve_stac_object().target for link in links]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcd8bb2c-67c0-40e0-ab91-9a5accb6bc71",
   "metadata": {},
   "source": [
    "We can plot the first item:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca94193d-41f9-4dd8-9601-48db4a71e639",
   "metadata": {},
   "outputs": [],
   "source": [
    "label_item = next(training_catalog.get_all_items())\n",
    "\n",
    "rasterio.plot.show(rasterio.open(label_item.assets[\"labels\"].href));"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48396c0e-aeb6-4645-b4e6-3511a2c076e1",
   "metadata": {},
   "source": [
    "Let's also load the names of each integer label code, which are common across all the items."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95d00da8-9dbf-4cea-9bb2-9e424517035e",
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = requests.get(label_item.assets[\"raster_values\"].href).json()\n",
    "labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b1bfac6-4530-4a6f-8b23-4215fd4129a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ipyleaflet import Map, GeoJSON\n",
    "\n",
    "data = {\n",
    "    \"type\": \"FeatureCollection\",\n",
    "    \"features\": [],\n",
    "}\n",
    "for label_item in label_items:\n",
    "    data[\"features\"].append(label_item.geometry)\n",
    "\n",
    "\n",
    "center = shapely.geometry.shape(label_item.geometry).centroid.bounds[:2][::-1]\n",
    "\n",
    "m = Map(center=(-32.5, 18.5), zoom=7)\n",
    "layer = GeoJSON(data=data)\n",
    "\n",
    "m.add_layer(layer)\n",
    "m.layout.max_width = \"600px\"\n",
    "\n",
    "m"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efc47c42-42a9-450a-a68d-af3df6241844",
   "metadata": {},
   "source": [
    "## Load the Sentinel-2 Scenes\n",
    "\n",
    "Just like before, we'll find a \"good\" (not too cloudy) Sentinel 2 scene using the Planetary Computer's STAC API. But remeber that we're dealing with a whole bunch of `label` scenes now, rather than just one. Because Sentinel scenes are so much larger that the `labels` chips, we'll do a single search rather than many small ones.\n",
    "\n",
    "First, let's verify that all the `labels` chips have the same datetime:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a45d371-72e9-424d-a985-55147d9dfe90",
   "metadata": {},
   "outputs": [],
   "source": [
    "# all the labels are from the same datetime\n",
    "set([label_item.datetime for label_item in label_items])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4327aa6-6f83-4430-850e-75d14b72349d",
   "metadata": {},
   "source": [
    "Now, let's compute a big bounding box that covers all the `labels` chips."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c989eb1-6cf1-4e83-b24d-546d70589409",
   "metadata": {},
   "outputs": [],
   "source": [
    "bboxes = np.array([label_item.bbox for label_item in label_items])\n",
    "mins = bboxes.min(0)\n",
    "maxes = bboxes.max(0)\n",
    "bbox = mins[0], mins[1], maxes[2], maxes[3]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef03a29d-cd87-410b-b661-ebf2eafeca6d",
   "metadata": {},
   "source": [
    "Find all the STAC items from `sentinel-2-l2a` within that bounding box for the month before and after the labels were created. Look back at our previous notebook if you need help."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2228b37d-004a-49f0-bb10-ff6b0ddfba06",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create the STAC client\n",
    "stac_client = pystac_client.Client.open(\n",
    "    \"https://planetarycomputer.microsoft.com/api/stac/v1/\"\n",
    ")\n",
    "\n",
    "date_range = \"2017-07-01/2017-09-01\"\n",
    "\n",
    "search = stac_client.search(\n",
    "    collections=[\"sentinel-2-l2a\"],\n",
    "    bbox=bbox,\n",
    "    datetime=date_range,\n",
    "    limit=500,\n",
    ")\n",
    "sentinel_items = list(search.get_all_items())\n",
    "len(sentinel_items)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79f62c91-5cc6-4240-b773-047717a38484",
   "metadata": {},
   "source": [
    "## Load the Data\n",
    "\n",
    "How best to get data from storage to the GPU is a large topic. We're going to take the easy way out and just load all the data into memory up front. This won't work if the entire dataset is larger than memory.\n",
    "\n",
    "We'll write a couple helper functions to\n",
    "\n",
    "1. Find the best Sentinel item for a given `labels` chip (\"best\" meaning the sentinel scene covers the chip, doesn't have many clouds)\n",
    "2. Load the `assets` from the best Sentinel item into a DataArray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2067dcf4-445d-4a97-80ae-8a2744a251ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_match(label_item, sentinel_items):\n",
    "    # make sure we pick a sentinel scene that overlaps substantially with the label\n",
    "    label_shape = shapely.geometry.shape(label_item.geometry)\n",
    "    items2 = [\n",
    "        item\n",
    "        for item in sentinel_items\n",
    "        if (\n",
    "            shapely.geometry.shape(item.geometry).intersection(label_shape).area\n",
    "            / label_shape.area\n",
    "        )\n",
    "        > 0.9\n",
    "    ]\n",
    "    sentinel_item = sorted(\n",
    "        items2, key=lambda item: pystac.extensions.eo.EOExtension.ext(item).cloud_cover\n",
    "    )[0]\n",
    "    return sentinel_item\n",
    "\n",
    "\n",
    "def get_item(label_item, sentinel_items, assets):\n",
    "    assets = list(assets)\n",
    "    labels = rioxarray.open_rasterio(\n",
    "        label_item.assets[\"labels\"].href,\n",
    "    ).squeeze()\n",
    "\n",
    "    sentinel_item = find_match(label_item, sentinel_items)\n",
    "    bounds = tuple(round(x, 0) for x in labels.rio.bounds())\n",
    "\n",
    "    data = (\n",
    "        stackstac.stack(\n",
    "            planetary_computer.sign(sentinel_item).to_dict(),\n",
    "            assets=assets,\n",
    "            dtype=\"float32\",\n",
    "            resolution=10,\n",
    "            bounds=bounds,\n",
    "            epsg=labels.rio.crs.to_epsg(),\n",
    "        )\n",
    "        .squeeze()\n",
    "        .compute(scheduler=\"single-threaded\")\n",
    "    )\n",
    "\n",
    "    assert data.shape[1:] == labels.shape\n",
    "    data = data.assign_coords(x=labels.x.data, y=labels.y.data)\n",
    "    data /= 4000\n",
    "    data = np.clip(data, 0, 1)\n",
    "    return data, labels.astype(\"int64\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a0fde33-920f-4ec4-84b9-d2f21af141d1",
   "metadata": {},
   "source": [
    "That takes a bit of time, so we'll do it in parallel using Dask. For this kind of ad-hoc parallelism, `dask.delayed` is a nice option."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a559132c-5d30-4ca4-a8ae-c29bf930030f",
   "metadata": {},
   "outputs": [],
   "source": [
    "get_item_ = dask.delayed(get_item, nout=2)\n",
    "assets = (\"B03\", \"B04\", \"B05\")\n",
    "\n",
    "Xys = [get_item_(label_item, sentinel_items, assets) for label_item in label_items]\n",
    "Xys = dask.persist(*Xys)\n",
    "\n",
    "Xys = dask.compute(*Xys)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df03bc80-07c8-4882-b079-b05db0d28dfe",
   "metadata": {},
   "source": [
    "We'll wrap that up in a small pytorch `Dataset` class, which can be passed to a pytorch `DataLoader` (see https://pytorch.org/tutorials/beginner/basics/data_tutorial.html) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26457e4f-8000-4f7a-ad43-17e52fcf33af",
   "metadata": {},
   "outputs": [],
   "source": [
    "class STACDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, data):\n",
    "        self.data = data\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        X, y = self.data[idx]\n",
    "        X = torch.as_tensor(X.data)\n",
    "        y = torch.as_tensor(y.data)\n",
    "        return X, y\n",
    "\n",
    "\n",
    "dataset = STACDataset(Xys)\n",
    "loader = torch.utils.data.DataLoader(dataset, batch_size=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f5d3e4c-41d7-4a0c-a79f-2794d67e234e",
   "metadata": {},
   "source": [
    "## Model Definition\n",
    "\n",
    "Now let's define our segmentation model (thanks to Caleb Robinson, a data scientist from Microsoft's AI for Good program, for help with the loader and the model)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d48504dd-85be-43cd-a341-6faae249ad4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import segmentation_models_pytorch as smp\n",
    "import torch\n",
    "\n",
    "model = smp.Unet(\n",
    "    encoder_name=\"resnet18\",\n",
    "    in_channels=len(assets),\n",
    "    classes=len(labels),\n",
    ")\n",
    "\n",
    "loss = smp.utils.losses.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(list(model.parameters()))\n",
    "metrics = [\n",
    "    smp.utils.metrics.IoU(threshold=0.5),\n",
    "]\n",
    "\n",
    "train_epoch = smp.utils.train.TrainEpoch(\n",
    "    model,\n",
    "    loss=loss,\n",
    "    metrics=metrics,\n",
    "    optimizer=optimizer,\n",
    "    device=device,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b10a3a3d-f7ec-446d-9763-909daccd3d1c",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "Let's start training the model. We'll train it for a single epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2c49929-eb4b-4e04-a4c2-37bb2bd791db",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "with device_ctx(device):\n",
    "    train_epoch.run(loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efb2511a-cb78-4930-b977-663a6448eafc",
   "metadata": {},
   "source": [
    "And let's visualize the output for the first 10 chips."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4f30780-f713-4ac7-9f06-685b8553073d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "def visualize(X, y):\n",
    "    with torch.no_grad():\n",
    "        X2 = (X.reshape((1,) + X.shape)).to(device)\n",
    "        output = model.predict(X2)\n",
    "\n",
    "    data = np.asarray(X.to(\"cpu\"))[0]  # just the first band\n",
    "    output = np.asarray(output[0].argmax(0).to(\"cpu\"))\n",
    "\n",
    "    fig, (ax1, ax2, ax3) = plt.subplots(ncols=3, sharey=True, sharex=True)\n",
    "\n",
    "    ax1.imshow(data, cmap=\"gray\")\n",
    "\n",
    "    rasterio.plot.show(np.asarray(y.to(\"cpu\")), ax=ax2)\n",
    "    rasterio.plot.show(output, ax=ax3)\n",
    "\n",
    "    return ax1, ax2, ax3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "124c3cea-1438-4a3d-a498-4879ca19fef7",
   "metadata": {},
   "source": [
    "The order is `(visible, actual, predicted)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00fb2e76-856f-481e-9d64-cc12401f88e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(10):\n",
    "    visualize(*dataset[i]);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4f47b68-5f87-439e-a7e6-fcea2bcc1ae0",
   "metadata": {},
   "source": [
    "Now let's train it a bit longer and see how well we do."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bced1e66-1224-4b6a-87b0-8a91819740a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "with device_ctx(0):\n",
    "    for i in range(40):\n",
    "        train_epoch.run(loader)\n",
    "        print(i, end=\"\\r\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "156eee71-136f-4fbd-98c1-4691b746f153",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(10):\n",
    "    visualize(*dataset[i]);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a37a89f-c4e4-4489-9f6a-208c1c693113",
   "metadata": {},
   "source": [
    "### Exercise: Play around!\n",
    "\n",
    "In the remaining time, I encourage you to play around with the data preprocessing and model architecture. See what you can do."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}