end2end_geospatial_ml_mapping.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8N04-Lg3OHnQ"
      },
      "source": [
        "# An Interactive Workflow for Geospatial ML Mapping.\n",
        "\n",
        "<img src=\"https://github.com/GeoAIAfrica/interactive_geospatial_mapping/blob/main/static/image_analysis.webp?raw=1\" width=\"50%\" />\n",
        "\n",
        "<a href=\"https://colab.research.google.com/github/GeoAIAfrica/interactive_geospatial_mapping/blob/main/end2end_geospatial_ml_mapping.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
        "\n",
        "**© AMLD 2024**. `MIT` License.\n",
        "\n",
        "**Authors:** [Akram Zaytar](https://www.linkedin.com/in/akramz/), [Gilles Q. Hacheme](https://www.linkedin.com/in/gilles-q-hacheme-a0956ab7/), [Aisha Alaagib](https://www.linkedin.com/in/aishaalaagib/), [Girmaw A. Tadesse](https://www.linkedin.com/in/girmaw-abebe-tadesse/).\n",
        "\n",
        "**Introduction:**\n",
        "\n",
        "In this notebook, we will present an end-to-end workflow for geospatial mapping using deep neural networks.\n",
        "\n",
        "We aim to cover the following:\n",
        "1. Introduction to Geospatial Data\n",
        "2. Pick a place & period of interest!\n",
        "3. Load imagery into the interactive map!\n",
        "4. Create a few labels for your object of interest! Export the Image/Mask!\n",
        "5. _Train!_  Data augmentation, Regularization, Fine-tuning, object vectorization!\n",
        "6. Export pixel-wise metrics for the local region!\n",
        "7. Run the model over a much bigger region & find other instances!\n",
        "8. Be ambitious! augmentation techniques, fusing layers (Sentinel-1), multi-class, other encoders/architectures, join our community!\n",
        "\n",
        "**Topics:**\n",
        "\n",
        "Content: <font color='blue'>`Geospatial Data Analysis`</font>, <font color='blue'>`Computer Vision`</font>, <font color='blue'>`Interactive Mapping`</font>.\n",
        "Level: <font color='grey'>`Beginner`</font>, <font color='grey'>`Intermediate`</font>\n",
        "\n",
        "**Outcome:**\n",
        "\n",
        "- *The basics* of Geosptial Data analysis: learn about data formats & types, foundational concepts.\n",
        "- *Interactive Mapping*: how can you acquire & prepare the inputs (satellite images) and create the targets (objects of interest) within a Jupyter notebook environment.\n",
        "- *Geospatial ML*: learn about geospatial Train/Val splitting, Data augmentation, Regularization, Fine-tuning, object vectorization, Evaluation, and inference.\n",
        "\n",
        "**Prerequisites:**\n",
        "- Basic to intermediate knowledge in Python and machine learning.\n",
        "- Familiarity with satellite imagery and geospatial data is beneficial but not mandatory.\n",
        "- Installation of necessary software and tools as detailed in the workshop's GitHub repository README file.\n",
        "- Participants are encouraged to install and set up the required tools prior to the workshop for a more efficient hands-on session.\n",
        "\n",
        "**Before you start:**\n",
        "\n",
        "For this practical, you will need to use a GPU to speed up training. To do this, go to the \"Runtime\" menu in Colab, select \"Change runtime type\" and then in the popup menu, choose \"GPU\" in the \"Hardware accelerator\" box."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hgBU3QMNOHnR"
      },
      "source": [
        "## Installation and Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ylgB0lHGOvTH"
      },
      "outputs": [],
      "source": [
        "%pip install rioxarray -q\n",
        "%pip install leafmap -q\n",
        "%pip install torchgeo==0.4.0 -q\n",
        "%pip install lightning -q\n",
        "%pip install geopandas -q\n",
        "%pip install localtileserver -q\n",
        "%pip install pystac_client -q\n",
        "%pip install planetary_computer -q\n",
        "%pip install stackstac -q"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hu_7AiubOHnS"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import warnings\n",
        "import shutil\n",
        "import hashlib\n",
        "import warnings\n",
        "import requests\n",
        "import subprocess\n",
        "from random import *\n",
        "from tqdm import tqdm\n",
        "from datetime import *\n",
        "from pathlib import Path\n",
        "from functools import reduce\n",
        "from functools import partial\n",
        "from typing import Any, Dict, cast\n",
        "\n",
        "warnings.filterwarnings(\"ignore\", category=FutureWarning)\n",
        "if os.getenv(\"COLAB_RELEASE_TAG\"):\n",
        "    from google.colab import output\n",
        "    output.enable_custom_widget_manager()\n",
        "\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import rasterio\n",
        "import stackstac\n",
        "import pystac_client\n",
        "import geopandas as gpd\n",
        "import planetary_computer\n",
        "from shapely.geometry import *\n",
        "from rasterio.features import rasterize\n",
        "from rasterio.transform import from_bounds\n",
        "from localtileserver import get_leaflet_tile_layer, TileClient\n",
        "from rasterio.windows import Window\n",
        "import leafmap\n",
        "from shapely.geometry import *\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from lightning import LightningDataModule, LightningModule\n",
        "from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint\n",
        "from lightning.pytorch.loggers import TensorBoardLogger\n",
        "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
        "from torch.utils.data import DataLoader\n",
        "from torchgeo.datasets import RasterDataset\n",
        "from torchgeo.datasets import stack_samples\n",
        "from torchgeo.samplers import RandomBatchGeoSampler, GridGeoSampler\n",
        "from torch import Tensor\n",
        "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
        "from lightning.pytorch import LightningModule\n",
        "import segmentation_models_pytorch as smp\n",
        "from torchgeo.transforms import AugmentationSequential\n",
        "import kornia.augmentation as K\n",
        "from torchmetrics import MetricCollection\n",
        "from torchmetrics.classification import (\n",
        "    MulticlassJaccardIndex,\n",
        "    Precision,\n",
        "    Recall,\n",
        "    F1Score,\n",
        ")\n",
        "import lightning.pytorch as pl\n",
        "from lightning.pytorch.callbacks import ModelCheckpoint, EarlyStopping\n",
        "from lightning.pytorch.loggers import TensorBoardLogger"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jbw7Owev6Q46"
      },
      "outputs": [],
      "source": [
        "# Create directories to be used later.\n",
        "store = Path(\"./data\")\n",
        "store.mkdir(exist_ok=True)\n",
        "logs = Path(\"./logs\")\n",
        "store.mkdir(exist_ok=True)\n",
        "results = Path(\"./results\")\n",
        "store.mkdir(exist_ok=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GZtoHPK4OHnS"
      },
      "source": [
        "---\n",
        "\n",
        "## Introduction to Geospatial Data!\n",
        "\n",
        "Geospatial data refers to **information that can be associated with locations on Earth**. It comes with attributes like **coordinates** and **geometry**.\n",
        "\n",
        "Examples of geospatial data:\n",
        "- Weather information.\n",
        "- Transportation networks.\n",
        "- Population density.\n",
        "\n",
        "There are two types of geospatial data:\n",
        "\n",
        "<div style=\"text-align:left;\"> <figure> <img width=\"500px\" src=\"https://i0.wp.com/pangeography.com/wp-content/uploads/2022/05/Raster_vector_tikz.png\" /> <figcaption style=\"font-size:small;\">Image credit: <a href=\"https://pangeography.com/geographic-data-structure-vector-data-and-raster-data/\">Pan Geography</a></figcaption> </figure> </div>\n",
        "\n",
        "- **Vector**: Points, Lines, Polygons, etc. Vector objects are geometries that may have multiple attributes. It is saved in a vector file (e.g., `Shapefile` (.shp), `GeoJSON`, among others).\n",
        "- **Raster**: represented as a grid of pixels, each pixel contains a value that represents a measurement. Raster data is stored in format like `GeoTIFF` and `NetCDF`.\n",
        "\n",
        "For **Vector** and **Raster** data, we need utilities that map a pixel's coordinates to a location on Earth. **Coordinate Reference Systems** (CRS) combine an earth model and a projection system, which translates the 3D Earth surface onto a 2D plane. Commonly used CRS include `WGS84`, often used for `GPS` data, and `UTM`, a set of projections that divide the world into a series of 6-degree longitudinal zones. When working with geospatial data, it is crucial to ensure that all datasets have the same CRS to avoid errors when aligning them.\n",
        "\n",
        "Converting between different CRS is known as \"reprojection.\" Care must be taken during reprojection to maintain data integrity, especially when working with large areas or when precision is crucial. Cooordinate reference systems (CRS) can take you from the geometric coordinates (numbers) to the earth's surface. `GeoPandas` allows us to inspect the CRS and reproject it if necessary.\n",
        "\n",
        "In `Python`, we can use the [rasterio](https://github.com/rasterio/rasterio) library to read/write raster data and the [geopandas](https://github.com/geopandas/geopandas) library for vector data.\n",
        "\n",
        "We present examples on how to read & visualize both vector & raster data:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "--ZXwv7aOHnS"
      },
      "outputs": [],
      "source": [
        "# Use geopandas API to read all countries in the world\n",
        "world = gpd.read_file(gpd.datasets.get_path(\"naturalearth_lowres\"))\n",
        "print(f\"Vector data coordinate reference system is — {world.crs}\")\n",
        "\n",
        "# Filter African countries\n",
        "africa = world[world[\"continent\"] == \"Africa\"]\n",
        "\n",
        "# Plot Africa\n",
        "fig, ax = plt.subplots(figsize=(5, 5))\n",
        "ax = africa.boundary.plot(ax=ax)\n",
        "ax.axis(\"off\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UIyJz0ScOHnT"
      },
      "source": [
        "We can use `leafmap` to interactively visualize a very high-resolution image (you can source others from [here](https://openaerialmap.org/)):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lH_1LBsUOHnT"
      },
      "outputs": [],
      "source": [
        "# Set the COG image URL\n",
        "img_url = \"https://oin-hotosm.s3.us-east-1.amazonaws.com/65de85797175970001f718f1/0/65de85797175970001f718f2.tif\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fCqlVRufOHnT"
      },
      "outputs": [],
      "source": [
        "# Create a tile server from local raster file\n",
        "client = TileClient(img_url)\n",
        "\n",
        "# Create ipyleaflet tile layer from that server\n",
        "t = get_leaflet_tile_layer(client)\n",
        "\n",
        "# Create the map\n",
        "m = leafmap.Map(center=client.center(), zoom=8)\n",
        "m.add(t)\n",
        "m"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wt-vxWODOHnT"
      },
      "source": [
        "We can also focus on a specific sub-region of the image and plot it:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "912ml9JLOHnT"
      },
      "outputs": [],
      "source": [
        "# Read a 1000x1000 window at the center of the image\n",
        "with rasterio.open(img_url) as src:\n",
        "    height, width = src.shape\n",
        "    hc, wc = int(height / 2), int(width / 2)\n",
        "    y0, x0 = hc - 1_000, wc - 1_000\n",
        "    arr = src.read(window=Window(y0, x0, 1_000, 1_000))\n",
        "\n",
        "# Check the array shape\n",
        "arr.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H7ht1EDNOHnT"
      },
      "outputs": [],
      "source": [
        "fig, ax = plt.subplots(figsize=(5, 5))\n",
        "ax.imshow(arr.transpose(1, 2, 0))\n",
        "ax.axis(\"off\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RvjM3Oqgh1DN"
      },
      "outputs": [],
      "source": [
        "# Clear the layers\n",
        "m.clear_layers()\n",
        "\n",
        "# No need for the previous map\n",
        "del [m]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WRNixF0KOHnT"
      },
      "source": [
        "---\n",
        "## 1. Pick a Region & Period of Interest\n",
        "\n",
        "After briefly going over geospatial data types, let's use `leafmap` to pick a region of interest:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VQoXl6QyOHnT"
      },
      "outputs": [],
      "source": [
        "# Initialize the Map object\n",
        "m = leafmap.Map()\n",
        "\n",
        "# Output the map object\n",
        "m"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dUD6YOP9OHnU"
      },
      "source": [
        "**Note: before running the next cell, zoom to a region of interest and draw a polygon that represents your area of interest (AOI)**.\n",
        "\n",
        "Let's save the ROI as a `GeoJSON` file:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7ZlaxMdkQfFK"
      },
      "outputs": [],
      "source": [
        "# Set this variable to false in case you don't want to label!\n",
        "labeling = False"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "leDeMYa5OHnU"
      },
      "outputs": [],
      "source": [
        "if labeling:\n",
        "\n",
        "    # Get the region of interest\n",
        "    roi = box(*m.user_roi_bounds())\n",
        "\n",
        "    # Save it as a `GeoJSON` file\n",
        "    gpd.GeoDataFrame(geometry=[roi], crs=\"EPSG:4326\").to_file(\"./data/roi.geojson\")\n",
        "\n",
        "else:\n",
        "\n",
        "    # Download the ROI file\n",
        "    !wget -nc -P data/ https://raw.githubusercontent.com/GeoAIAfrica/interactive_geospatial_mapping/main/data/roi.geojson"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PLHWuCnFOHnU"
      },
      "source": [
        "After setting the region of interest, let's set the period of interest:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m5H4ILIzOHnU"
      },
      "outputs": [],
      "source": [
        "start_date = \"2023-01-01\"\n",
        "end_date = \"2024-01-01\"\n",
        "start_date, end_date"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pQVcluF8OHnU"
      },
      "source": [
        "---\n",
        "\n",
        "## 2. Load imagery into the interactive map!\n",
        "\n",
        "After setting the region and period of interest, we will use [Microsoft Planetary Computer](https://planetarycomputer.microsoft.com/catalog) to do the following:\n",
        "1. Search for cloud-free Sentinel-2 images that correspond to the ROI/period.\n",
        "2. Stack the found items on the time dimension and create a `rioxarray` object.\n",
        "3. Crop the images to the region of interest and calculate the median, finally, save the resulting Mosaic."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kYpGQEw5OHnU"
      },
      "outputs": [],
      "source": [
        "# Get the catalog\n",
        "api_url = \"https://planetarycomputer.microsoft.com/api/stac/v1\"\n",
        "catalog = pystac_client.Client.open(api_url, modifier=planetary_computer.sign_inplace)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7q_Lt_PWOHnU"
      },
      "outputs": [],
      "source": [
        "# Set the search parameters\n",
        "product = \"sentinel-2-l2a\"\n",
        "roi = gpd.read_file(\"./data/roi.geojson\").geometry.values[0]\n",
        "bbox = roi.bounds\n",
        "max_cloud_percent = 0\n",
        "period = f\"{start_date}/{end_date}\"\n",
        "print(f\"- Bounding box: {bbox}\")\n",
        "print(f\"- Satellite product: {product}\")\n",
        "print(f\"- Period: {period}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xacLnl0_OHnU"
      },
      "outputs": [],
      "source": [
        "# Search!\n",
        "results = catalog.search(\n",
        "    collections=[product],\n",
        "    bbox=bbox,\n",
        "    datetime=period,\n",
        "    query={\n",
        "        \"eo:cloud_cover\": {\"lte\": max_cloud_percent}\n",
        "    },  # cloud_cover less than or equal to 0%\n",
        ")\n",
        "\n",
        "# Get the items\n",
        "items = results.get_all_items()\n",
        "print(f\"Items found: {len(items)}.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ug9Ic4bKOHnU"
      },
      "source": [
        "We can now stack the items and create a dataset object:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0_jwtjtIOHnU"
      },
      "outputs": [],
      "source": [
        "images = stackstac.stack(items, assets=[\"B02\", \"B03\", \"B04\", \"B08\"])\n",
        "images"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4etthH-tOHnU"
      },
      "source": [
        "Let's create a Mosaic focused on our region of interest and save/visualize it:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iD2lx1erOHnU"
      },
      "outputs": [],
      "source": [
        "# Crop to the region of interest\n",
        "X = images.rio.clip([roi], crs=f\"EPSG:4326\", drop=True)\n",
        "\n",
        "# Create the Mosaic\n",
        "X = X.median(dim=\"time\", keep_attrs=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P3MrCziAOHnU"
      },
      "outputs": [],
      "source": [
        "# Save the ROI image\n",
        "X.rio.write_nodata(0.0, encoded=True, inplace=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "arkZwhVuOHnU"
      },
      "outputs": [],
      "source": [
        "X.rio.to_raster(\"./data/X.tif\", compress=\"lzw\", dtype=\"float32\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MmF_SiRAOHnU"
      },
      "outputs": [],
      "source": [
        "# Clear the previous drawing of the region of interest\n",
        "if labeling:\n",
        "    m.clear_drawings()\n",
        "\n",
        "# Add the image to the map\n",
        "m.add_raster(\"./data/X.tif\", bands=[3, 2, 1], layer_name=\"X\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9vrpzpO0OHnU"
      },
      "source": [
        "---\n",
        "\n",
        "## 3. Create a few labels!\n",
        "\n",
        "**Note: go back to the map widget and draw polygons for the object of interest! When you are done, continue from here...**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s6U0jnrqCklJ"
      },
      "outputs": [],
      "source": [
        "def contribute(file_path):\n",
        "    \"\"\"\n",
        "    Uploads a file to a predefined URL using a POST request.\n",
        "\n",
        "    Args:\n",
        "        file_path (str): The path to the file to be uploaded.\n",
        "\n",
        "    Returns:\n",
        "        A requests.Response object containing the server's response to the HTTP request.\n",
        "    \"\"\"\n",
        "    url = \"https://labeluploader.azurewebsites.net/upload\"\n",
        "    files = {'file': (file_path, open(file_path, 'rb'), 'application/geo+json')}\n",
        "\n",
        "    try:\n",
        "        response = requests.post(url, files=files)\n",
        "        return response\n",
        "    except requests.exceptions.RequestException as e:\n",
        "        print(f\"An error occurred: {e}\")\n",
        "        return None"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TV1NnMwkOHnU"
      },
      "outputs": [],
      "source": [
        "if labeling:\n",
        "\n",
        "    # Get the geometries from the map\n",
        "    geoms = [Polygon(e[\"geometry\"][\"coordinates\"][0]) for e in m.draw_features]\n",
        "\n",
        "    # Set the name of the object you just labeled!\n",
        "    object_name = \"\"\n",
        "    assert object_name, \"Please set the name of the object you labeled!\"\n",
        "\n",
        "    # Get the current time as a `datetime` object\n",
        "    now = datetime.now()\n",
        "\n",
        "    # Create a `GeoDataFrame` from the geometries (EPSG:4326), the object name, and the date\n",
        "    gdf = gpd.GeoDataFrame(geometry=geoms, crs=\"EPSG:4326\")\n",
        "    gdf[\"object\"] = object_name\n",
        "    gdf[\"time\"] = now\n",
        "\n",
        "    # Generate a random string hash for the file name `{hash}.geojson`\n",
        "    fn_stem = hashlib.md5(f\"{object_name}{now.strftime('%Y-%m-%d-%H-%M-%S')}\".encode()).hexdigest()\n",
        "    file_path = Path(f\"./data/{fn_stem}.geojson\")\n",
        "\n",
        "    # Save the file\n",
        "    gdf.to_file(file_path)\n",
        "\n",
        "    # Contribute your labels to GeoAI Africa\n",
        "    contribute(file_path)\n",
        "\n",
        "else:\n",
        "\n",
        "    # We download the labels\n",
        "    !wget -nc -P data/ https://raw.githubusercontent.com/GeoAIAfrica/interactive_geospatial_mapping/main/data/labels.geojson"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YlDjjr8COHnU"
      },
      "source": [
        "Now, let's create a function that takes the labels' `GeoDataFrame` and exports the mask file:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ueWfFNyZOHnU"
      },
      "outputs": [],
      "source": [
        "def create_mask(labels, crs, resolution, mask_path, pos_class=1, neg_class=2):\n",
        "    \"\"\"\n",
        "    Create a raster mask from vector labels.\n",
        "\n",
        "    Args:\n",
        "        labels (GeoDataFrame): A GeoDataFrame containing geometries (labels).\n",
        "        crs (CRS): Coordinate Reference System to use for the output raster.\n",
        "        resolution (float): The pixel size in the units of the CRS.\n",
        "        mask_path (str): Path where the raster mask will be saved.\n",
        "        pos_class (int, optional): Value to assign for positive class (default: 1).\n",
        "        neg_class (int, optional): Value to assign for negative class (unused, default: 0).\n",
        "\n",
        "    The function converts the geometries in the GeoDataFrame to a raster mask.\n",
        "    Each pixel in the mask represents whether it falls inside a geometry (positive class)\n",
        "    or outside (negative class, typically 0). The output is a single-band GeoTIFF file.\n",
        "    \"\"\"\n",
        "\n",
        "    # Ensure the labels' CRS matches the target CRS\n",
        "    if str(labels.crs) != str(crs):\n",
        "        labels = labels.to_crs(crs)\n",
        "\n",
        "    # Get geometries from the labels and create a list of values for rasterization\n",
        "    geoms = labels.geometry.tolist()\n",
        "    vals = [pos_class] * len(geoms)\n",
        "\n",
        "    # Compute bounds for the output mask and calculate the transform\n",
        "    minx, miny, maxx, maxy = labels.unary_union.bounds\n",
        "    mask_transform = from_bounds(\n",
        "        minx,\n",
        "        miny,\n",
        "        maxx,\n",
        "        maxy,\n",
        "        width=int((maxx - minx) / resolution),\n",
        "        height=int((maxy - miny) / resolution),\n",
        "    )\n",
        "\n",
        "    # Create metadata for the output raster file\n",
        "    mask_metadata = {\n",
        "        \"driver\": \"GTiff\",  # File format\n",
        "        \"dtype\": \"uint8\",  # Data type of the raster\n",
        "        \"nodata\": None,  # NoData value; None implies no NoData value\n",
        "        \"width\": int((maxx - minx) / resolution),  # Raster width in pixels\n",
        "        \"height\": int((maxy - miny) / resolution),  # Raster height in pixels\n",
        "        \"count\": 1,  # Number of bands in the raster\n",
        "        \"crs\": crs,  # Coordinate Reference System\n",
        "        \"transform\": mask_transform,  # Affine transformation parameters\n",
        "        \"compress\": \"lzw\",  # Compression algorithm\n",
        "        \"predictor\": 2,  # Predictor for compression\n",
        "    }\n",
        "\n",
        "    # Write the raster mask to a file\n",
        "    with rasterio.open(mask_path, \"w\", **mask_metadata) as out_img:\n",
        "        # Rasterize the geometries into an array\n",
        "        mask_arr = rasterize(\n",
        "            tuple(zip(geoms, vals)),\n",
        "            out_shape=(mask_metadata[\"height\"], mask_metadata[\"width\"]),\n",
        "            transform=mask_transform,\n",
        "            fill=neg_class,  # Fill value for negative class\n",
        "            default_value=neg_class,  # Default fill value\n",
        "        )\n",
        "        # Write the array to the raster band\n",
        "        out_img.write_band(1, mask_arr)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "APxDdx4GOHnU"
      },
      "source": [
        "We can now export a mask that corresponds to the labels that we just created:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8OwmDoUHOHnV"
      },
      "outputs": [],
      "source": [
        "# We want the mask to have the same CRS as the image\n",
        "img_crs = rasterio.open(\"./data/X.tif\").crs\n",
        "\n",
        "# We know that the resolution is 10 meters (Sentinel-2)\n",
        "resolution = 10\n",
        "\n",
        "# We load the labels\n",
        "labels = gpd.read_file(\"./data/labels.geojson\").to_crs(img_crs)\n",
        "\n",
        "# We set the path of the mask to export\n",
        "mask_path = Path(\"./data/mask.tif\")\n",
        "\n",
        "# Write!\n",
        "create_mask(labels, img_crs, resolution, mask_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WKfqg7U9OHnV"
      },
      "outputs": [],
      "source": [
        "# Visualize it alongside the image\n",
        "m.add_raster(\"./data/mask.tif\", bands=[1], layer_name=\"y\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1H1ihBVfOHnV"
      },
      "source": [
        "---\n",
        "\n",
        "## Train!\n",
        "\n",
        "In this section, we will use [`TorchGeo`](https://github.com/microsoft/torchgeo) to train a `UNet` model with a `ResNet18` encoder over the labeled pixels."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h-m-7MqPOHnV"
      },
      "source": [
        "We will start by creating a `dataset` class that can index into a raster image:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9TJ0DzvsOHnV"
      },
      "outputs": [],
      "source": [
        "class SingleRasterDataset(RasterDataset):\n",
        "    \"\"\"\n",
        "    A class representing a single raster dataset, inheriting from RasterDataset.\n",
        "\n",
        "    This class is designed to handle individual raster datasets by specifying a file name.\n",
        "    It sets up the dataset by extracting the directory of the given file as its root directory\n",
        "    and allows for the application of transformations to the dataset.\n",
        "\n",
        "    Attributes:\n",
        "        filename_regex (str): The base name of the file specified for the dataset.\n",
        "                              This attribute is intended for internal use to identify\n",
        "                              the dataset file.\n",
        "\n",
        "    Parameters:\n",
        "        fn (str): The path to the single raster file. This path is used to extract the\n",
        "                  file name and the directory for initializing the dataset.\n",
        "        transforms (callable, optional): A function/transform that takes in a sample and returns\n",
        "                                         a transformed version. These transforms are applied to\n",
        "                                         the dataset items. Default is None.\n",
        "\n",
        "    Note:\n",
        "        The `transforms` parameter allows for preprocessing or data augmentation operations\n",
        "        to be applied to the dataset. Ensure that any transforms provided are compatible\n",
        "        with raster data.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, fn, transforms=None):\n",
        "        self.filename_regex = os.path.basename(fn)\n",
        "        # Initialize the base RasterDataset class with the directory of the file and any transforms\n",
        "        super().__init__(root=os.path.dirname(fn), transforms=transforms)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oGZUtIrXOHnV"
      },
      "outputs": [],
      "source": [
        "def preprocess(sample, remove_bbox=True, max_val=7938, bands=[2, 1, 0]):\n",
        "    \"\"\"\n",
        "    Preprocesses a given sample by applying band selection, normalization, and optional removal of bounding boxes.\n",
        "\n",
        "    This function modifies the input sample in-place by selecting specified bands from the \"image\" field,\n",
        "    normalizing the selected image bands by a specified maximum value, converting the image data type to float,\n",
        "    and optionally removing the bounding box information. If present, the \"mask\" field is squeezed to remove\n",
        "    singleton dimensions and converted to long data type.\n",
        "\n",
        "    Parameters:\n",
        "        sample (dict): A dictionary representing a sample from a dataset. The sample is expected to contain\n",
        "                       an \"image\" key with image data and optionally \"mask\" and \"bbox\" keys for the segmentation\n",
        "                       mask and bounding box information, respectively.\n",
        "        remove_bbox (bool, optional): A flag to indicate whether bounding box information (\"bbox\" key) should be\n",
        "                                      removed from the sample. Defaults to True.\n",
        "        max_val (int, optional): The maximum value used for normalizing the image data. Defaults to 7938, which\n",
        "                                 is often used for satellite imagery normalization.\n",
        "        bands (list of int, optional): The indices of the bands to be selected from the image. Defaults to [2, 1, 0],\n",
        "                                       typically corresponding to the RGB bands of satellite imagery.\n",
        "\n",
        "    Returns:\n",
        "        dict: The preprocessed sample with the image data selected, normalized, and converted to float type,\n",
        "              the mask (if present) squeezed and converted to long type, and the bounding box information (if present\n",
        "              and `remove_bbox` is True) removed.\n",
        "\n",
        "    Note:\n",
        "        The function modifies the input `sample` dictionary in-place, but also returns the modified dictionary\n",
        "        for convenience and chaining operations.\n",
        "    \"\"\"\n",
        "    if \"image\" in sample:\n",
        "        # Select specified bands and normalize the image\n",
        "        sample[\"image\"] = sample[\"image\"][bands]\n",
        "        sample[\"image\"] = (sample[\"image\"] / max_val).float()\n",
        "    if \"mask\" in sample:\n",
        "        # Squeeze the mask to remove singleton dimensions and convert to long data type\n",
        "        sample[\"mask\"] = sample[\"mask\"].squeeze().long()\n",
        "    if remove_bbox and \"bbox\" in sample:\n",
        "        # Remove the bounding box information if specified\n",
        "        del sample[\"bbox\"]\n",
        "    return sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rHsJOBPqOHnb"
      },
      "source": [
        ".. and now we create the data module class, responsible for creating **datalaoders** used for *batch generation*:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CjohfA4mOHnb"
      },
      "outputs": [],
      "source": [
        "class SegmentationDataModule(LightningDataModule):\n",
        "    \"\"\"PyTorch Lightning DataModule for a segmentation task.\"\"\"\n",
        "\n",
        "    def __init__(\n",
        "        self,\n",
        "        img_path: Path,\n",
        "        mask_path: Path,\n",
        "        batch_size: int = 64,\n",
        "        patch_size: int = 256,\n",
        "        batches_per_epoch: int = 512,\n",
        "        workers: int = 4,\n",
        "    ):\n",
        "        \"\"\"\n",
        "        Initialize the SegmentationDataModule.\n",
        "        Args:\n",
        "        img_path (Path): The filepath to the input image file.\n",
        "        mask_path (Path): The filepath to the mask image file.\n",
        "        batch_size (int, optional): The number of samples per batch during training. Defaults to 64.\n",
        "        patch_size (int, optional): The size of patches to be extracted from the images. Defaults to 256.\n",
        "        batches_per_epoch (int, optional): The number of batches per training epoch. Defaults to 512.\n",
        "        workers (int, optional): The number of worker threads for data loading. Defaults to 4.\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "\n",
        "        # Verify that the image file exists\n",
        "        if not img_path.exists():\n",
        "            raise FileNotFoundError(\"The image file does not exist.\")\n",
        "        if not mask_path.exists():\n",
        "            raise FileNotFoundError(\"The mask file does not exist.\")\n",
        "\n",
        "        # Save the path to the input file\n",
        "        self.X_file = img_path\n",
        "        self.y_file = mask_path\n",
        "\n",
        "        # Save the rest of the hyperparameters\n",
        "        self.batch_size = batch_size\n",
        "        self.patch_size = patch_size\n",
        "        self.batches_per_epoch = batches_per_epoch\n",
        "        self.workers = workers\n",
        "\n",
        "        self.train_ds = None\n",
        "        self.val_ds = None\n",
        "\n",
        "    def setup(self, stage=None):\n",
        "        \"\"\"\n",
        "        Setup method to prepare the datasets for training and validation.\n",
        "        This method calculates class weights if they have not been calculated and creates the\n",
        "        datasets if they have not been created.\n",
        "        Args:\n",
        "        stage (str, optional): The stage for which the setup is being run ('fit' or 'test'). Defaults to None.\n",
        "        \"\"\"\n",
        "\n",
        "        # Only setup if datasets are not already initialized\n",
        "        if self.train_ds is None:\n",
        "            # Create the training dataset\n",
        "            self.train_img_ds = SingleRasterDataset(self.X_file, transforms=preprocess)\n",
        "            self.train_mask_ds = SingleRasterDataset(self.y_file, transforms=preprocess)\n",
        "            self.train_mask_ds.is_image = False\n",
        "            self.train_ds = self.train_img_ds & self.train_mask_ds\n",
        "\n",
        "            # Because of lack of labels, we will use the same mask for validation\n",
        "            self.val_img_ds = SingleRasterDataset(self.X_file, transforms=preprocess)\n",
        "            self.val_mask_ds = SingleRasterDataset(self.y_file, transforms=preprocess)\n",
        "            self.val_mask_ds.is_image = False\n",
        "            self.val_ds = self.val_img_ds & self.val_mask_ds\n",
        "\n",
        "    def train_dataloader(self):\n",
        "        \"\"\"\n",
        "        Prepare the dataloader for the training dataset.\n",
        "        Returns:\n",
        "        DataLoader: Dataloader for the training dataset.\n",
        "        \"\"\"\n",
        "        sampler = RandomBatchGeoSampler(\n",
        "            self.train_ds,\n",
        "            size=self.patch_size,\n",
        "            batch_size=self.batch_size,\n",
        "            length=self.batches_per_epoch * self.batch_size,\n",
        "        )\n",
        "        return DataLoader(\n",
        "            self.train_ds,\n",
        "            batch_sampler=sampler,\n",
        "            num_workers=self.workers,\n",
        "            collate_fn=stack_samples,\n",
        "        )\n",
        "\n",
        "    def val_dataloader(self):\n",
        "        \"\"\"\n",
        "        Prepare the dataloader for the validation dataset.\n",
        "        Returns:\n",
        "        DataLoader: Dataloader for the validation dataset.\n",
        "        \"\"\"\n",
        "        sampler = RandomBatchGeoSampler(\n",
        "            self.val_ds,\n",
        "            size=self.patch_size,\n",
        "            batch_size=self.batch_size,\n",
        "            length=self.batches_per_epoch * self.batch_size,\n",
        "        )\n",
        "        return DataLoader(\n",
        "            self.train_ds,\n",
        "            batch_sampler=sampler,\n",
        "            num_workers=self.workers,\n",
        "            collate_fn=stack_samples,\n",
        "        )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nx2p8P9_OHnc"
      },
      "source": [
        "Let's test the data module class:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EKht1fXAOHnc"
      },
      "outputs": [],
      "source": [
        "store = Path(\"./data\")\n",
        "img = store / \"X.tif\"\n",
        "mask = store / \"mask.tif\"\n",
        "assert img.exists()\n",
        "assert mask.exists()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "em4lzffCOHnc"
      },
      "outputs": [],
      "source": [
        "dm = SegmentationDataModule(\n",
        "    img_path=img, mask_path=mask, batch_size=8, patch_size=128, workers=0\n",
        ")\n",
        "dm.setup()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qB-0K1_EOHnc"
      },
      "source": [
        "Let's visualize a few samples to make sure the data loaders are working correctly:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MJbxphayOHnc"
      },
      "outputs": [],
      "source": [
        "# Create the training data loader\n",
        "train_dl = dm.train_dataloader()\n",
        "batch = next(iter(train_dl))\n",
        "imgs = batch[\"image\"]\n",
        "masks = batch[\"mask\"]\n",
        "\n",
        "# Visualize the images and masks side-by-side using matplotlib\n",
        "fig, axs = plt.subplots(nrows=2, ncols=8, figsize=(32, 8))\n",
        "for i, (img, mask) in enumerate(zip(imgs, masks)):\n",
        "    axs[0, i].imshow(img.numpy().transpose(1, 2, 0))\n",
        "    axs[0, i].axis(\"off\")\n",
        "    axs[1, i].imshow(mask.numpy())\n",
        "    axs[1, i].axis(\"off\")\n",
        "plt.tight_layout()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FiR229tyOHnc"
      },
      "source": [
        "Let's create the training class addapted from [`TorchGeo`](https://github.com/microsoft/torchgeo/blob/e04e1a53fd6a21506693d53f8a8519dbf4261817/torchgeo/trainers/segmentation.py#L24):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_VzNDMo6OHnc"
      },
      "outputs": [],
      "source": [
        "class SemanticSegmentationTask(LightningModule):\n",
        "\n",
        "    def __init__(self, **kwargs: Any) -> None:\n",
        "        \"\"\"\n",
        "        Adapted from `TorchGeo`.\n",
        "        \"\"\"\n",
        "\n",
        "        # Call the superclass constructor\n",
        "        super().__init__()\n",
        "\n",
        "        # Init the loss and model choices\n",
        "        self.loss = None\n",
        "        self.model = None\n",
        "        self.class_weights = None\n",
        "\n",
        "        # Creates `self.hparams` from kwargs\n",
        "        self.save_hyperparameters()\n",
        "        self.hyperparams = cast(Dict[str, Any], self.hparams)\n",
        "\n",
        "        # Validate the hyperparameters\n",
        "        self._validate_ignore_index(kwargs[\"ignore_index\"])\n",
        "        self._validate_model(kwargs[\"model\"])\n",
        "        self._validate_loss(kwargs[\"loss\"])\n",
        "\n",
        "        # Get the number of classes\n",
        "        self.num_classes = self.hyperparams[\"num_classes\"]\n",
        "\n",
        "        # Set the indices for each class\n",
        "        self.class2idx = dict()\n",
        "        for i, class_name in enumerate(self.hyperparams[\"class_names\"]):\n",
        "            self.class2idx[class_name] = i + 1\n",
        "\n",
        "        # Set the ignore index\n",
        "        self.ignore_index = kwargs[\"ignore_index\"]\n",
        "\n",
        "        # Call the config task method\n",
        "        self.config_task()\n",
        "\n",
        "        # Set the color jitter parameters\n",
        "        color_jitter_params = {\n",
        "            \"brightness\": 0.2,\n",
        "            \"contrast\": 0.2,\n",
        "            \"saturation\": 0.2,\n",
        "            \"hue\": 0.1,\n",
        "            \"p\": 0.8,\n",
        "        }\n",
        "\n",
        "        # If any augmentation is enabled, create the augmentation pipeline\n",
        "        augmentations = list()\n",
        "        augmentations.append(K.RandomHorizontalFlip(p=0.5))\n",
        "        augmentations.append(K.RandomVerticalFlip(p=0.5))\n",
        "        augmentations.append(K.RandomRotation(degrees=90.0, p=0.5))\n",
        "        augmentations.append(K.ColorJitter(**color_jitter_params))\n",
        "\n",
        "        # Create the augmentation function\n",
        "        self.aug = AugmentationSequential(*augmentations, data_keys=[\"image\", \"mask\"])\n",
        "\n",
        "        # Set the metrics of interest\n",
        "        metrics = {\n",
        "            \"JaccardIndex\": MulticlassJaccardIndex(\n",
        "                num_classes=self.hyperparams[\"num_classes\"],\n",
        "                average=None,\n",
        "                ignore_index=self.ignore_index,\n",
        "            ),\n",
        "            \"Precision\": Precision(\n",
        "                task=\"multiclass\",\n",
        "                num_classes=self.hyperparams[\"num_classes\"],\n",
        "                average=None,\n",
        "                ignore_index=self.ignore_index,\n",
        "            ),\n",
        "            \"Recall\": Recall(\n",
        "                task=\"multiclass\",\n",
        "                num_classes=self.hyperparams[\"num_classes\"],\n",
        "                average=None,\n",
        "                ignore_index=self.ignore_index,\n",
        "            ),\n",
        "            \"F1\": F1Score(\n",
        "                task=\"multiclass\",\n",
        "                num_classes=self.hyperparams[\"num_classes\"],\n",
        "                average=None,\n",
        "                ignore_index=self.ignore_index,\n",
        "            ),\n",
        "        }\n",
        "        self.train_metrics = MetricCollection(metrics, prefix=\"train_\")\n",
        "        self.val_metrics = self.train_metrics.clone(prefix=\"val_\")\n",
        "\n",
        "        # Intend to save the validation losses\n",
        "        self.batch_val_losses = []\n",
        "\n",
        "    def _validate_ignore_index(self, ignore_index):\n",
        "        if not isinstance(ignore_index, (int, type(None))):\n",
        "            raise ValueError(\"ignore_index must be an int or None\")\n",
        "        if (ignore_index is not None) and (self.hyperparams[\"loss\"] == \"jaccard\"):\n",
        "            warnings.warn(\n",
        "                \"ignore_index has no effect on training when loss='jaccard'\",\n",
        "                UserWarning,\n",
        "            )\n",
        "\n",
        "    def _validate_model(self, model):\n",
        "        valid_models = [\n",
        "            \"unet\",\n",
        "            \"deeplabv3+\",\n",
        "            \"unet++\",\n",
        "            \"manet\",\n",
        "            \"linknet\",\n",
        "            \"fpn\",\n",
        "            \"pspnet\",\n",
        "            \"pan\",\n",
        "            \"deeplabv3\",\n",
        "        ]\n",
        "        if model not in valid_models:\n",
        "            raise ValueError(\n",
        "                f\"Model type '{model}' is not valid. \"\n",
        "                f\"Currently, only supports {valid_models}.\"\n",
        "            )\n",
        "\n",
        "    def _validate_loss(self, loss):\n",
        "        valid_losses = [\"ce\", \"jaccard\", \"focal\"]\n",
        "        if loss not in valid_losses:\n",
        "            raise ValueError(\n",
        "                f\"Loss type '{loss}' is not valid. \"\n",
        "                f\"Currently, supports {valid_losses} loss.\"\n",
        "            )\n",
        "\n",
        "    def config_task(self) -> None:\n",
        "        \"\"\"Configures the task based on kwargs parameters passed to the constructor.\"\"\"\n",
        "        self._init_model()\n",
        "        self._init_loss()\n",
        "\n",
        "    def _init_model(self):\n",
        "        if self.hyperparams[\"model\"] == \"unet\":\n",
        "            self.model = smp.Unet(\n",
        "                encoder_name=self.hyperparams[\"backbone\"],\n",
        "                encoder_weights=self.hyperparams[\"weights\"],\n",
        "                in_channels=self.hyperparams[\"in_channels\"],\n",
        "                classes=self.hyperparams[\"num_classes\"],\n",
        "            )\n",
        "        elif self.hyperparams[\"model\"] == \"deeplabv3+\":\n",
        "            self.model = smp.DeepLabV3Plus(\n",
        "                encoder_name=self.hyperparams[\"backbone\"],\n",
        "                encoder_weights=self.hyperparams[\"weights\"],\n",
        "                in_channels=self.hyperparams[\"in_channels\"],\n",
        "                classes=self.hyperparams[\"num_classes\"],\n",
        "            )\n",
        "        elif self.hyperparams[\"model\"] == \"unet++\":\n",
        "            self.model = smp.UnetPlusPlus(\n",
        "                encoder_name=self.hyperparams[\"backbone\"],\n",
        "                encoder_depth=5,\n",
        "                encoder_weights=self.hyperparams[\"weights\"],\n",
        "                in_channels=self.hyperparams[\"in_channels\"],\n",
        "                classes=self.hyperparams[\"num_classes\"],\n",
        "            )\n",
        "        elif self.hyperparams[\"model\"] == \"manet\":\n",
        "            self.model = smp.MAnet(\n",
        "                encoder_name=self.hyperparams[\"backbone\"],\n",
        "                encoder_depth=5,\n",
        "                encoder_weights=self.hyperparams[\"weights\"],\n",
        "                in_channels=self.hyperparams[\"in_channels\"],\n",
        "                classes=self.hyperparams[\"num_classes\"],\n",
        "            )\n",
        "        elif self.hyperparams[\"model\"] == \"linknet\":\n",
        "            self.model = smp.Linknet(\n",
        "                encoder_name=self.hyperparams[\"backbone\"],\n",
        "                encoder_depth=5,\n",
        "                encoder_weights=self.hyperparams[\"weights\"],\n",
        "                in_channels=self.hyperparams[\"in_channels\"],\n",
        "                classes=self.hyperparams[\"num_classes\"],\n",
        "            )\n",
        "        elif self.hyperparams[\"model\"] == \"fpn\":\n",
        "            self.model = smp.FPN(\n",
        "                encoder_name=self.hyperparams[\"backbone\"],\n",
        "                encoder_depth=5,\n",
        "                encoder_weights=self.hyperparams[\"weights\"],\n",
        "                in_channels=self.hyperparams[\"in_channels\"],\n",
        "                classes=self.hyperparams[\"num_classes\"],\n",
        "            )\n",
        "        elif self.hyperparams[\"model\"] == \"pspnet\":\n",
        "            self.model = smp.PSPNet(\n",
        "                encoder_name=self.hyperparams[\"backbone\"],\n",
        "                encoder_weights=self.hyperparams[\"weights\"],\n",
        "                in_channels=self.hyperparams[\"in_channels\"],\n",
        "                classes=self.hyperparams[\"num_classes\"],\n",
        "            )\n",
        "        elif self.hyperparams[\"model\"] == \"pan\":\n",
        "            self.model = smp.PAN(\n",
        "                encoder_name=self.hyperparams[\"backbone\"],\n",
        "                encoder_weights=self.hyperparams[\"weights\"],\n",
        "                in_channels=self.hyperparams[\"in_channels\"],\n",
        "                classes=self.hyperparams[\"num_classes\"],\n",
        "            )\n",
        "        elif self.hyperparams[\"model\"] == \"deeplabv3\":\n",
        "            self.model = smp.DeepLabV3(\n",
        "                encoder_name=self.hyperparams[\"backbone\"],\n",
        "                encoder_depth=5,\n",
        "                encoder_weights=self.hyperparams[\"weights\"],\n",
        "                in_channels=self.hyperparams[\"in_channels\"],\n",
        "                classes=self.hyperparams[\"num_classes\"],\n",
        "            )\n",
        "\n",
        "    def _init_loss(self):\n",
        "        if self.hyperparams[\"loss\"] == \"ce\":\n",
        "            if self.class_weights is not None:\n",
        "                self.class_weights = self.class_weights.to(self.device)\n",
        "            self.loss = nn.CrossEntropyLoss(\n",
        "                weight=self.class_weights,\n",
        "                ignore_index=-1000 if self.ignore_index is None else self.ignore_index,\n",
        "            )\n",
        "        elif self.hyperparams[\"loss\"] == \"jaccard\":\n",
        "            self.loss = smp.losses.JaccardLoss(\n",
        "                mode=\"multiclass\", classes=self.hyperparams[\"num_classes\"]\n",
        "            )\n",
        "        elif self.hyperparams[\"loss\"] == \"focal\":\n",
        "            self.loss = smp.losses.FocalLoss(\n",
        "                \"multiclass\", ignore_index=self.ignore_index, normalized=True\n",
        "            )\n",
        "\n",
        "    def forward(self, *args: Any, **kwargs: Any) -> Any:\n",
        "        \"\"\"Forward pass of the model.\"\"\"\n",
        "        return self.model(*args, **kwargs)\n",
        "\n",
        "    def training_step(self, *args: Any, **kwargs: Any) -> Tensor:\n",
        "        \"\"\"Compute and return the training loss.\"\"\"\n",
        "        batch = self.aug(args[0])  # if self.do_augment else args[0]\n",
        "        x = batch[\"image\"]\n",
        "        y = batch[\"mask\"]\n",
        "        y_hat = self(x)\n",
        "\n",
        "        # Report loss\n",
        "        y_hat_hard = y_hat.argmax(dim=1)\n",
        "        loss = self.loss(y_hat, y)\n",
        "        self.log(\"train_loss\", loss, on_step=True, on_epoch=False)\n",
        "\n",
        "        # Report metrics\n",
        "        self.train_metrics(y_hat_hard, y)\n",
        "        return cast(Tensor, loss)\n",
        "\n",
        "    def validation_step(self, *args: Any, **kwargs: Any) -> None:\n",
        "        \"\"\"Compute validation loss and log example predictions.\"\"\"\n",
        "        batch = args[0]\n",
        "        x = batch[\"image\"]\n",
        "        y = batch[\"mask\"]\n",
        "        y_hat = self(x)\n",
        "\n",
        "        # Report loss\n",
        "        y_hat_hard = y_hat.argmax(dim=1)\n",
        "        loss = self.loss(y_hat, y)\n",
        "        self.log(\"val_loss\", loss, on_step=False, on_epoch=True)\n",
        "        self.batch_val_losses.append(loss)\n",
        "\n",
        "        # Report metrics\n",
        "        self.val_metrics(y_hat_hard, y)\n",
        "\n",
        "    def on_train_epoch_end(self) -> None:\n",
        "        \"\"\"Logs epoch level training metrics.\"\"\"\n",
        "        train_metrics = self.train_metrics.compute()\n",
        "        new_metrics = dict()\n",
        "        for k in train_metrics.keys():\n",
        "            for category, cat_idx in self.class2idx.items():\n",
        "                new_metrics[f\"{k}_{category}\"] = train_metrics[k][cat_idx]\n",
        "        self.log_dict(new_metrics)\n",
        "        self.train_metrics.reset()\n",
        "\n",
        "    def on_validation_epoch_end(self) -> None:\n",
        "        \"\"\"Logs epoch level validation metrics.\"\"\"\n",
        "\n",
        "        # Calculate the rest of the metrics\n",
        "        val_metrics = self.val_metrics.compute()\n",
        "        new_metrics = dict()\n",
        "        for k in val_metrics.keys():\n",
        "            for category, cat_idx in self.class2idx.items():\n",
        "                new_metrics[f\"{k}_{category}\"] = val_metrics[k][cat_idx]\n",
        "\n",
        "        # Estimate the validation loss\n",
        "        val_batch_losses = self.batch_val_losses\n",
        "        val_loss = torch.nanmean(torch.stack(val_batch_losses))\n",
        "        new_metrics[\"val_loss\"] = val_loss\n",
        "\n",
        "        self.log_dict(new_metrics)\n",
        "        self.val_metrics.reset()\n",
        "        self.batch_val_losses = list()\n",
        "\n",
        "    def predict_step(self, *args: Any, **kwargs: Any) -> Tensor:\n",
        "        \"\"\"Compute and return the predictions.\"\"\"\n",
        "        batch = args[0]\n",
        "        x = batch[\"image\"]\n",
        "        y_hat: Tensor = self(x).softmax(dim=1)\n",
        "        return y_hat\n",
        "\n",
        "    def configure_optimizers(self) -> Dict[str, Any]:\n",
        "        \"\"\"\n",
        "        Configure the optimizer and learning rate scheduler based on the hyperparameters.\n",
        "        Returns:\n",
        "            A dictionary containing the optimizer and learning rate scheduler.\n",
        "        \"\"\"\n",
        "\n",
        "        # Retrieve the optimizer name, learning rate, and weight decay from the hyperparameters\n",
        "        optimizer_name = self.hyperparams[\"optimizer_name\"]\n",
        "        learning_rate = self.hyperparams[\"learning_rate\"]\n",
        "        weight_decay = self.hyperparams[\"weight_decay\"]\n",
        "\n",
        "        # Select the optimizer based on the specified name\n",
        "        if optimizer_name == \"SGD\":\n",
        "            optimizer = torch.optim.SGD(\n",
        "                self.model.parameters(), lr=learning_rate, weight_decay=weight_decay\n",
        "            )\n",
        "        elif optimizer_name == \"Adam\":\n",
        "            optimizer = torch.optim.Adam(\n",
        "                self.model.parameters(), lr=learning_rate, weight_decay=weight_decay\n",
        "            )\n",
        "        elif optimizer_name == \"RMSProp\":\n",
        "            optimizer = torch.optim.RMSprop(\n",
        "                self.model.parameters(), lr=learning_rate, weight_decay=weight_decay\n",
        "            )\n",
        "        elif optimizer_name == \"AdamW\":\n",
        "            optimizer = torch.optim.AdamW(\n",
        "                self.model.parameters(), lr=learning_rate, weight_decay=weight_decay\n",
        "            )\n",
        "        else:\n",
        "            raise ValueError(f\"Unknown optimizer: {optimizer_name}\")\n",
        "\n",
        "        # Retrieve the scheduler name from the hyperparameters\n",
        "        scheduler_name = self.hyperparams[\"scheduler_name\"]\n",
        "\n",
        "        # Select the learning rate scheduler based on the specified name\n",
        "        if scheduler_name == \"ReduceLROnPlateau\":\n",
        "            scheduler = ReduceLROnPlateau(\n",
        "                optimizer, patience=self.hyperparams[\"learning_rate_schedule_patience\"]\n",
        "            )\n",
        "        elif scheduler_name == \"CosineAnnealingLR\":\n",
        "            scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n",
        "                optimizer, T_max=self.hyperparams[\"T_max\"]\n",
        "            )\n",
        "        elif scheduler_name == \"StepLR\":\n",
        "            scheduler = torch.optim.lr_scheduler.StepLR(\n",
        "                optimizer,\n",
        "                step_size=self.hyperparams[\"step_size\"],\n",
        "                gamma=self.hyperparams[\"gamma\"],\n",
        "            )\n",
        "        else:\n",
        "            raise ValueError(f\"Unknown scheduler: {scheduler_name}\")\n",
        "\n",
        "        # Return the optimizer and learning rate scheduler\n",
        "        return {\n",
        "            \"optimizer\": optimizer,\n",
        "            \"lr_scheduler\": {\n",
        "                \"scheduler\": scheduler,\n",
        "                \"monitor\": \"val_loss\",\n",
        "            },\n",
        "        }"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nkOXhsDvOHnc"
      },
      "source": [
        "We will fix a set of hyperparameters:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "byB0qfGaOHnc"
      },
      "outputs": [],
      "source": [
        "# Define experiment setup and hyperparameters\n",
        "\n",
        "# Check if there is a GPU\n",
        "is_gpu = torch.cuda.is_available()\n",
        "\n",
        "# Name of the experiment for tracking\n",
        "experiment_name = \"urban_mappings\"\n",
        "\n",
        "# Loss function to be used (categorical crossentropy in this case)\n",
        "loss = \"ce\"\n",
        "\n",
        "# Size of batches for training\n",
        "batch_size = 128  # 256\n",
        "\n",
        "# Size of the patches to be extracted from the images\n",
        "patch_size = 128\n",
        "\n",
        "# Number of batches to process in an epoch\n",
        "batches_per_epoch = 32\n",
        "\n",
        "# Number of worker processes for loading data\n",
        "workers = 4 if is_gpu else 0\n",
        "\n",
        "# Number of epochs with no improvement after which training will be stopped\n",
        "early_stopping_patience = 10\n",
        "\n",
        "# Minimum/Maximum number of epochs to run before early stopping\n",
        "min_epochs = 1  # 5\n",
        "max_epochs = 10  # 50\n",
        "\n",
        "# Number of classes in the dataset\n",
        "num_classes = 3\n",
        "\n",
        "# Architecture/Encoder of the model\n",
        "arch = \"unet\"\n",
        "backbone = \"mobilenet_v2\"\n",
        "\n",
        "# Pre-trained weights to initialize the backbone (Imagenet weights)\n",
        "weights = \"imagenet\"\n",
        "\n",
        "# Initial learning rate for training\n",
        "learning_rate = 0.0001\n",
        "\n",
        "# Number of epochs with no improvement on validation loss after which learning rate will be reduced\n",
        "lr_schedule_patience = 10\n",
        "\n",
        "# Weight decay (L2 penalty) for regularization\n",
        "weight_decay = 1e-2\n",
        "\n",
        "# Scheduler for adjusting learning rate\n",
        "scheduler_name = \"ReduceLROnPlateau\"\n",
        "\n",
        "# Optimizer for training\n",
        "optimizer_name = \"AdamW\"\n",
        "\n",
        "# Index to be ignored in loss computation, useful for masked areas in segmentation tasks\n",
        "ignore_index = 0\n",
        "\n",
        "# Paths to the input image and mask files, ensuring both exist\n",
        "img_path = store / \"X.tif\"\n",
        "mask_path = store / \"mask.tif\"\n",
        "assert img_path.exists() and mask_path.exists()  # Ensure both paths exist"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7f1tqRCPOHnc"
      },
      "source": [
        "We can train!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wd9JspaIOHnc"
      },
      "outputs": [],
      "source": [
        "# Set the experiment name and directory\n",
        "results_dir = Path(\"./results\")\n",
        "results_dir.mkdir(exist_ok=True)\n",
        "experiment_dir = results_dir / experiment_name\n",
        "experiment_dir"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "USCEn5TVOHnc"
      },
      "outputs": [],
      "source": [
        "# Create the data module\n",
        "dm = SegmentationDataModule(\n",
        "    img_path=img_path,\n",
        "    mask_path=mask_path,\n",
        "    batches_per_epoch=batches_per_epoch,\n",
        "    batch_size=batch_size,\n",
        "    patch_size=patch_size,\n",
        "    workers=workers,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5-C5T0N8OHnc"
      },
      "outputs": [],
      "source": [
        "# Create the task\n",
        "task = SemanticSegmentationTask(\n",
        "    model=arch,\n",
        "    backbone=backbone,\n",
        "    weights=weights,\n",
        "    in_channels=3,\n",
        "    num_classes=num_classes,\n",
        "    ignore_index=ignore_index,\n",
        "    learning_rate=learning_rate,\n",
        "    learning_rate_schedule_patience=lr_schedule_patience,\n",
        "    loss=loss,\n",
        "    weight_decay=weight_decay,\n",
        "    scheduler_name=scheduler_name,\n",
        "    optimizer_name=optimizer_name,\n",
        "    class_names=[\"urban\", \"background\"],\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AcBCK16dOHnc"
      },
      "outputs": [],
      "source": [
        "# Create the checkpoint and early stopping callbacks\n",
        "checkpoint_callback = ModelCheckpoint(\n",
        "    monitor=\"val_loss\",\n",
        "    dirpath=experiment_dir,\n",
        "    save_top_k=1,\n",
        "    mode=\"min\",\n",
        ")\n",
        "\n",
        "# Create the early stopping callback\n",
        "early_stop_callback = EarlyStopping(\n",
        "    monitor=\"val_loss\", patience=early_stopping_patience, verbose=True, mode=\"min\"\n",
        ")\n",
        "\n",
        "# Create the TensorBoard logger\n",
        "tb_logger = TensorBoardLogger(save_dir=\"logs/\", name=experiment_name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-C1RbntSOHnc"
      },
      "outputs": [],
      "source": [
        "# Trainer definition\n",
        "trainer = pl.Trainer(\n",
        "    logger=[tb_logger],\n",
        "    max_epochs=max_epochs,\n",
        "    min_epochs=min_epochs,\n",
        "    callbacks=[checkpoint_callback, early_stop_callback],\n",
        "    precision=16,\n",
        "    accelerator=\"gpu\" if is_gpu else \"cpu\",\n",
        "    devices=[0] if is_gpu else \"auto\",\n",
        ")\n",
        "trainer.fit(model=task, datamodule=dm)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1SnptYQ9OHnc"
      },
      "source": [
        "Let's report the local Jaccard, F1, recall, and precision metrics:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qyWAuJn-OHnd"
      },
      "outputs": [],
      "source": [
        "report = dict()\n",
        "class_of_interest = \"urban\"\n",
        "for metric in [\"F1\", \"JaccardIndex\", \"Recall\", \"Precision\"]:\n",
        "    report[metric] = round(\n",
        "        float(trainer.logged_metrics[f\"val_{metric}_{class_of_interest}\"]), 3\n",
        "    )\n",
        "report"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xYUsNKLxOHnd"
      },
      "source": [
        "---\n",
        "\n",
        "## Inference\n",
        "\n",
        "In this section, we will use our trained model to predict over the original image, vectorize, postprocess, and save our predictions:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YjBhtRoGOHnd"
      },
      "outputs": [],
      "source": [
        "def predict(\n",
        "    img_path,\n",
        "    model_path,\n",
        "    out_fp,\n",
        "    patch_size=128,\n",
        "    batch_size=1,\n",
        "    padding=1,\n",
        "    num_workers=4,\n",
        "    gpu=0,\n",
        "):\n",
        "    \"\"\"Segments a given image using a fine-tuned model and saves the output.\n",
        "\n",
        "    This function performs semantic segmentation on an input image using a specified model checkpoint. It processes\n",
        "    the image in patches, predicts each patch's segmentation mask, and stitches them together to form a complete output\n",
        "    mask, which is then saved to a specified output file path.\n",
        "\n",
        "    Parameters:\n",
        "        img_path (str): Path to the input image file.\n",
        "        model_path (str): Path to the pre-trained model checkpoint.\n",
        "        out_fp (str): Output file path where the prediction result will be saved.\n",
        "        patch_size (int, optional): Size of the patches to be extracted from the images. Defaults to 128.\n",
        "        batch_size (int, optional): Number of samples per batch during prediction. Defaults to 1.\n",
        "        padding (int, optional): Number of pixels to pad the patches. Defaults to 1.\n",
        "        num_workers (int, optional): Number of worker processes for data loading. Defaults to 4.\n",
        "        gpu (int, optional): The ID of the GPU to use for prediction. Defaults to 0.\n",
        "\n",
        "    Process:\n",
        "        1. Sets up device for prediction based on available GPU.\n",
        "        2. Loads the model from the checkpoint and prepares it for evaluation.\n",
        "        3. Creates a dataset and dataloader for the input image, dividing it into patches for efficient processing.\n",
        "        4. Iterates over the image patches, performs prediction, and accumulates the results.\n",
        "        5. Saves the aggregated predictions as a geospatial raster image, preserving the input image's spatial reference.\n",
        "\n",
        "    Note:\n",
        "        - The function assumes the use of a `SemanticSegmentationTask` model architecture for prediction.\n",
        "        - The input image is processed in patches to manage memory usage and adapt to different input sizes.\n",
        "        - The output raster will have a single band, float32 data type, and will be compressed using LZW compression.\n",
        "        - Padding is used to reduce edge effects in patch-based prediction, and is removed in the final output.\n",
        "        - The function uses a deterministic sampling strategy (GridGeoSampler) to ensure complete coverage of the input image.\n",
        "    \"\"\"\n",
        "\n",
        "    # Set the stride\n",
        "    stride = patch_size - padding * 2\n",
        "\n",
        "    # Set the device\n",
        "    device = torch.device(\n",
        "        f\"cuda:{gpu}\" if (gpu is not None) and torch.cuda.is_available() else \"cpu\"\n",
        "    )\n",
        "\n",
        "    # Load task and data\n",
        "    task = SemanticSegmentationTask.load_from_checkpoint(model_path)\n",
        "    task.freeze()\n",
        "    model = task.model\n",
        "    model = model.eval().to(device)\n",
        "\n",
        "    # Create a dataset object from a single image file\n",
        "    val_ds = SingleRasterDataset(\n",
        "        img_path, transforms=partial(preprocess, remove_bbox=False)\n",
        "    )\n",
        "\n",
        "    # Create the sampler (not random because we want to predict the whole image deterministically)\n",
        "    sampler = GridGeoSampler(val_ds, size=patch_size, stride=stride)\n",
        "\n",
        "    # Create the dataloader\n",
        "    val_dl = DataLoader(\n",
        "        val_ds,\n",
        "        sampler=sampler,\n",
        "        batch_size=batch_size,\n",
        "        num_workers=num_workers,\n",
        "        collate_fn=stack_samples,\n",
        "    )\n",
        "\n",
        "    # Open the input file\n",
        "    with rasterio.open(img_path) as f:\n",
        "        input_height, input_width = f.shape\n",
        "        profile = f.profile\n",
        "        transform = profile[\"transform\"]\n",
        "\n",
        "    # Initialize the output numpy array to zeros\n",
        "    output = np.zeros((input_height, input_width), dtype=np.float16)\n",
        "\n",
        "    # Create the enumerated to iterate over the TIF patches in batches\n",
        "    dl_enumerator = tqdm(val_dl)\n",
        "\n",
        "    # Iterate over the image batches are predict\n",
        "    for batch in dl_enumerator:\n",
        "        # Get the images and their bounding boxes\n",
        "        images = batch[\"image\"].to(device)\n",
        "        bboxes = batch[\"bbox\"]\n",
        "\n",
        "        with torch.inference_mode():\n",
        "            # Predict over all the images\n",
        "            y_hat = model(images)\n",
        "\n",
        "            # Get the predicted probabilities for the boma class\n",
        "            y_hat_boma = y_hat.softmax(dim=1)[:, 1, ...].cpu().numpy()\n",
        "\n",
        "        for i in range(len(bboxes)):\n",
        "            bb = bboxes[i]\n",
        "\n",
        "            left, top = ~transform * (bb.minx, bb.maxy)\n",
        "            right, bottom = ~transform * (bb.maxx, bb.miny)\n",
        "            left, right, top, bottom = (\n",
        "                int(np.round(left)),\n",
        "                int(np.round(right)),\n",
        "                int(np.round(top)),\n",
        "                int(np.round(bottom)),\n",
        "            )\n",
        "\n",
        "            assert right - left == patch_size\n",
        "            assert bottom - top == patch_size\n",
        "\n",
        "            output[\n",
        "                top + padding : bottom - padding, left + padding : right - padding\n",
        "            ] = y_hat_boma[i][padding:-padding, padding:-padding]\n",
        "\n",
        "        # Save predictions\n",
        "        profile[\"driver\"] = \"GTiff\"\n",
        "        profile[\"count\"] = 1\n",
        "        profile[\"dtype\"] = \"float32\"\n",
        "        profile[\"compress\"] = \"lzw\"\n",
        "        profile[\"predictor\"] = 2\n",
        "        profile[\"nodata\"] = 0\n",
        "        profile[\"blockxsize\"] = 512\n",
        "        profile[\"blockysize\"] = 512\n",
        "        profile[\"tiled\"] = True\n",
        "        profile[\"interleave\"] = \"pixel\"\n",
        "\n",
        "        # Save the file\n",
        "        with rasterio.open(out_fp, \"w\", **profile) as f:\n",
        "            f.write(output, 1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u1N0enXMOHnd"
      },
      "outputs": [],
      "source": [
        "# Get the best model checkpoint\n",
        "best_model_path = checkpoint_callback.best_model_path\n",
        "\n",
        "# Set the input/output paths\n",
        "img_path = Path(\"./data/X.tif\")\n",
        "out_fp = Path(\"./data/y_hat.tif\")\n",
        "\n",
        "# Run the main function\n",
        "predict(img_path, best_model_path, out_fp)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Hp8VN6FpOHnd"
      },
      "source": [
        "Let's vectorize our predictions and export object-wise metrics:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FQ1IPkThOHnd"
      },
      "outputs": [],
      "source": [
        "def raster_to_gdf(tif_path):\n",
        "    \"\"\"\n",
        "    Converts a raster file to a GeoDataFrame by polygonizing it.\n",
        "\n",
        "    This function takes the path to a raster (.tif) file, uses GDAL's gdal_polygonize utility to convert raster pixels\n",
        "    into polygons, and then loads these polygons into a GeoDataFrame. It is particularly useful for converting rasterized\n",
        "    masks or classifications into vector data for further geospatial analysis or visualization. The function specifically\n",
        "    filters for polygons corresponding to the class of interest (with DN value of 1), removes unnecessary columns, and\n",
        "    returns a clean GeoDataFrame.\n",
        "\n",
        "    Parameters:\n",
        "        tif_path (str): Path to the input raster (.tif) file.\n",
        "\n",
        "    Returns:\n",
        "        gpd.GeoDataFrame: A GeoDataFrame containing polygons for the specified class of interest from the raster. If the\n",
        "        raster conversion or loading fails, or if a timeout occurs, an empty GeoDataFrame is returned.\n",
        "    \"\"\"\n",
        "\n",
        "    # Ensure GDAL's gdal_polygonize.py is available\n",
        "    if not shutil.which(\"gdal_polygonize.py\"):\n",
        "        raise EnvironmentError(\n",
        "            \"gdal_polygonize.py is not available in the system path.\"\n",
        "        )\n",
        "\n",
        "    # Generate a random output shapefile name\n",
        "    output_shapefile = f\"output_{randint(0, 1e6)}.shp\"\n",
        "\n",
        "    # Construct the command\n",
        "    cmd = [\"gdal_polygonize.py\", tif_path, \"-f\", \"ESRI Shapefile\", output_shapefile]\n",
        "\n",
        "    try:\n",
        "        # Run the command with a timeout of 1 minute\n",
        "        subprocess.run(cmd, check=True, timeout=60)\n",
        "\n",
        "        # Load the shapefile into a GeoDataFrame\n",
        "        if not os.path.exists(output_shapefile):\n",
        "            raise FileNotFoundError(f\"Output shapefile not found: {output_shapefile}\")\n",
        "        gdf = gpd.read_file(output_shapefile)\n",
        "\n",
        "        # Clean up shapefile components\n",
        "        for ext in [\".shp\", \".shx\", \".dbf\", \".prj\"]:\n",
        "            os.remove(output_shapefile.replace(\".shp\", ext))\n",
        "\n",
        "        # Filter for the class of interest & return\n",
        "        return gdf[gdf[\"DN\"] == 1].drop(\"DN\", axis=1)\n",
        "\n",
        "    except subprocess.TimeoutExpired:\n",
        "\n",
        "        # Clean up shapefile components\n",
        "        for ext in [\".shp\", \".shx\", \".dbf\", \".prj\"]:\n",
        "            os.remove(output_shapefile.replace(\".shp\", ext))\n",
        "\n",
        "        # Return an empty GeoDataFrame in case of a timeout\n",
        "        return gpd.GeoDataFrame()\n",
        "\n",
        "    except Exception as e:\n",
        "        # Handle other potential exceptions\n",
        "        raise e"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "88myJQ0pOHnd"
      },
      "outputs": [],
      "source": [
        "# Vectorize the predicted masks\n",
        "gdf = raster_to_gdf(out_fp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fQ9XX6IKOHnd"
      },
      "outputs": [],
      "source": [
        "# Visualize the vectorized predictions\n",
        "ax = gdf.plot(color=\"red\")\n",
        "_ = ax.axis(\"off\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dy7l2OS5OHnd"
      },
      "source": [
        "Now, we postprocess the predictions in the following manner:\n",
        "\n",
        "1. Removes geometries that have a very small area.\n",
        "2. Simplifies the geometries.\n",
        "3. Simplifies the geometries by Dilation + Erosion.\n",
        "4. Fills any holes in the geometries."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7ZH13OmGOHnd"
      },
      "outputs": [],
      "source": [
        "# Filter valid geometries\n",
        "gdf = gdf[gdf.geometry.is_valid]\n",
        "\n",
        "# Filter small geometries\n",
        "q01_area = gdf.geometry.area.quantile(0.01)\n",
        "gdf = gdf[gdf.geometry.area > q01_area]\n",
        "\n",
        "# Simplify the geometries\n",
        "gdf[\"geometry\"] = gdf.geometry.simplify(10)\n",
        "\n",
        "# Dilate and erode the geometries\n",
        "buffer = 10\n",
        "gdf[\"geometry\"] = gdf.geometry.buffer(buffer)  # Dilation\n",
        "gdf[\"geometry\"] = gdf.geometry.buffer(-buffer)  # Erosion\n",
        "\n",
        "\n",
        "def fillit(row):\n",
        "    \"\"\"A function to fill holes below an area threshold in a polygon\"\"\"\n",
        "    newgeom = None\n",
        "    rings = [i for i in row[\"geometry\"].interiors]\n",
        "    if len(rings) > 0:  # If there are any rings\n",
        "        to_fill = [Polygon(ring) for ring in rings]\n",
        "        if len(to_fill) > 0:\n",
        "            newgeom = reduce(\n",
        "                lambda geom1, geom2: geom1.union(geom2), [row[\"geometry\"]] + to_fill\n",
        "            )  # Union the original geometry with all holes\n",
        "    if newgeom:\n",
        "        return newgeom\n",
        "    else:\n",
        "        return row[\"geometry\"]\n",
        "\n",
        "\n",
        "# Apply the function\n",
        "gdf[\"geometry\"] = gdf.apply(fillit, axis=1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ciC1-eS3OHnd"
      },
      "outputs": [],
      "source": [
        "# Visualize the final predictions\n",
        "ax = gdf.plot(color=\"red\")\n",
        "_ = ax.axis(\"off\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zUZqJdv4OHnd"
      },
      "source": [
        "Let's save our predictions and add them to the map!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hkA4PM846qDA"
      },
      "outputs": [],
      "source": [
        "# Save the predictions\n",
        "gdf.to_file(\"./data/predictions.geojson\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Rr1lcOakOHnd"
      },
      "outputs": [],
      "source": [
        "# Add to the map\n",
        "m.add_vector(\"./data/predictions.geojson\", layer_name=\"Predictions\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3hl6FZCqOHnd"
      },
      "source": [
        "---\n",
        "\n",
        "## Explore!\n",
        "\n",
        "While this tutorial covered essential techniques and approaches for end-to-end mapping, the field is vast, and there's much more to explore. Here are some directions you might consider to enhance your models' training and inference capabilities further:\n",
        "\n",
        "- **Different Architectures & Encoders**: Consider experimenting with different neural network architectures and encoders to find the optimal combination for your task. Select other architectures/encoders from [here](https://kornia.readthedocs.io/en/latest/augmentation.html).\n",
        "- **Advanced Data Augmentation Techniques**: Data augmentation is a powerful strategy to increase your training dataset, leading to more robust models. Investigate other augmentation techniques that could simulate more varied conditions or introduce more complex transformations. Look [here](https://kornia.readthedocs.io/en/latest/augmentation.html).\n",
        "- **Semi-Supervised Learning Techniques**: Semi-supervised learning can be particularly beneficial in scenarios where labeled data is scarce but unlabeled data is abundant. Explore how incorporating semi-supervised learning techniques can leverage unlabeled data to improve your model's performance.\n",
        "- **Fuse Data from Different Sources**: Combining data from different sensors, such as `Sentinel-1` (radar) and `Sentinel-2` (optical), can provide complementary information that enhances model understanding and performance.\n",
        "- **Scale to Multi-Class Categorization**: If your current model focuses on binary classification or a limited number of classes, consider expanding its capabilities to multi-class categorization. This expansion can increase the model's applicability and challenge it to capture more complex patterns in the data."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2bzx3YJmOHnd"
      },
      "source": [
        "---\n",
        "\n",
        "## Conclusion\n",
        "\n",
        "Throughout this notebook, we have ventured into the intricate domain of geospatial machine learning, highlighting the challenge of initiating machine learning-assisted geospatial mapping in the absence of labels, a task notably demanding due to the extensive steps involved. These steps range from acquiring satellite imagery, preparing and processing this imagery, establishing a labeling environment and budget, preparing a computational environment, to training the model. Such comprehensive requirements significantly hinder the application of deep learning models, particularly in regions with the most need, like Africa.\n",
        "\n",
        "We demonstrated the use of **open-source software**, **publicly available satellite imagery**, and **free computational resources** to conduct end-to-end mapping of a region of interest. This approach is crucial in low-resource settings, underscoring our objective to demonstrate an end-to-end workflow for mapping objects in satellite imagery utilizing publicly accessible resources at minimal costs. The ultimate goal is to empower geospatial machine learning applications across the continent, offering a beacon of hope for regions that stand to benefit the most from these advancements.\n",
        "\n",
        "As we conclude, it's important to remember that our journey through geospatial machine learning is just beginning. The field is ripe with opportunities for further exploration and innovation, promising to bring significant contributions to the world. Let's continue to learn, explore, and contribute to making a meaningful impact through geospatial machine learning. Happy learning!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fF28ubmROHnd"
      },
      "source": [
        "---\n",
        "\n",
        "## Resources\n",
        "\n",
        "### Tutorials\n",
        "\n",
        "- [Geospatial Primer](https://github.com/Akramz/geospatial-primer).\n",
        "- [Deep Learning Indaba Geospatial Tutorial](https://github.com/deep-learning-indaba/indaba-pracs-2023).\n",
        "- [Introduction to Geospatial Data](https://colab.research.google.com/drive/1-85h5tEB0AJYT8xQ5H1wtSnCafXuLTHo#scrollTo=JDT5jUmCiTH-).\n",
        "- [Geospatial Data Analysis](https://colab.research.google.com/drive/1Yfkm63OV3eCtR3IVB-4owi2DJgj2Wd84).\n",
        "- [Geospatial Deep learning: Getting started with TorchGeo](https://pytorch.org/blog/geospatial-deep-learning-with-torchgeo/).\n",
        "- [Automating GIS-processes Course]((https://autogis-site.readthedocs.io/en/latest/))\n",
        "- [Geospatial Data with Python: Shapely and Fiona](https://macwright.com/2012/10/31/gis-with-python-shapely-fiona.html)\n",
        "- [Introduction to Raster Data Processing in Open Source Python](https://www.earthdatascience.org/courses/use-data-open-source-python/intro-raster-data-python/raster-data-processing/).\n",
        "- [XArray fundamental](https://rabernat.github.io/research_computing_2018/xarray.html).\n",
        "- [XArray tutorials](https://github.com/xarray-contrib/xarray-tutorial).\n",
        "- [Visualization: contextily tutorial](https://geopandas.org/en/stable/gallery/plotting_basemap_background.html).\n",
        "\n",
        "\n",
        "### Libraries\n",
        "\n",
        "- [Shapely](https://github.com/shapely/shapely).\n",
        "- [GeoPandas](https://github.com/geopandas/geopandas).\n",
        "- [Contextily](https://github.com/geopandas/contextily).\n",
        "- [Rasterio](https://github.com/rasterio/rasterio).\n",
        "- [Xarray](https://github.com/pydata/xarray).\n",
        "- [RioXarray](https://github.com/corteva/rioxarray).\n",
        "- [TorchGeo](https://github.com/microsoft/torchgeo)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FiQJLBfTOHnd"
      },
      "source": [
        "---"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "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.10.13"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}