{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f77143aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (c) TorchGeo Contributors. All rights reserved.\n",
    "# Licensed under the MIT License."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3508e814",
   "metadata": {},
   "source": [
    "# NAIP Road Segmentation\n",
    "\n",
    "_Written by: Isaac Corley_\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "478038fe",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "The objective of this tutorial is to run inference with TorchGeo's pre-trained road segmentation model on NAIP imagery.\n",
    "\n",
    "This walkthrough covers:\n",
    "\n",
    "* How to stream a sample NAIP tile and road mask directly from remote GeoTIFF assets;\n",
    "* How to load `Unet_Weights.NAIP_RGBN_RESNET18_CHESAPEAKERSC` and its preprocessing pipeline;\n",
    "* How to run forward inference and compare predicted road masks against ground truth.\n",
    "\n",
    "Unlike a training tutorial, this notebook focuses only on model loading, preprocessing, and inference.\n",
    "\n",
    "### References\n",
    "\n",
    "[*Seeing the roads through the trees: A benchmark for modeling spatial dependencies with aerial imagery*, Robinson et al (2024)](https://arxiv.org/abs/2401.06762)\n",
    "\n",
    "ChesapeakeRSC dataset and implementation resources: [Hugging Face dataset card](https://huggingface.co/datasets/torchgeo/ChesapeakeRSC) | [GitHub source code](https://github.com/isaaccorley/ChesapeakeRSC)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d9c2c12",
   "metadata": {},
   "source": [
    "## Environment\n",
    "\n",
    "For the environment, install TorchGeo. Matplotlib is already available as a core dependency in the TorchGeo stack used by this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "711c651b",
   "metadata": {},
   "outputs": [],
   "source": [
    "!uv pip install torchgeo"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bf949da",
   "metadata": {},
   "source": [
    "## Imports\n",
    "\n",
    "We import:\n",
    "\n",
    "* `rasterio` to read remote GeoTIFFs;\n",
    "* `torch` and `numpy` for tensor preparation;\n",
    "* `matplotlib` for visualization;\n",
    "* TorchGeo model weights + model constructor for inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d58ccb5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import rasterio\n",
    "import torch\n",
    "\n",
    "from torchgeo.models import Unet_Weights, unet"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0160ad02",
   "metadata": {},
   "source": [
    "## Dataset\n",
    "\n",
    "We use one NAIP image tile and one corresponding road mask from the ChesapeakeRSC assets.\n",
    "\n",
    "To keep the tutorial lightweight, both files are read directly from remote HTTPS URLs in ``rasterio`` (no local download step)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a4384be",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_url = 'https://github.com/isaaccorley/ChesapeakeRSC/raw/66b29f478ee5263f51a8620eb3a34093a970523d/assets/8782_image.tif'\n",
    "mask_url = 'https://github.com/isaaccorley/ChesapeakeRSC/raw/66b29f478ee5263f51a8620eb3a34093a970523d/assets/8782_mask.tif'\n",
    "\n",
    "with rasterio.open(image_url) as src:\n",
    "    image_np = src.read().astype(np.float32)\n",
    "image = torch.from_numpy(image_np).unsqueeze(0)\n",
    "\n",
    "with rasterio.open(mask_url) as src:\n",
    "    mask = torch.from_numpy(src.read(1))\n",
    "    target = (mask == 9) | (mask == 12)  # road=9, road-occluded-by-tree-canopy=12\n",
    "    target = target.to(torch.uint8)\n",
    "\n",
    "rgb = image[0, :3].permute(1, 2, 0).byte()\n",
    "image.shape, image.dtype, target.shape, target.dtype, rgb.shape, rgb.dtype"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab5ba1e8",
   "metadata": {},
   "source": [
    "## Visualize Input Image\n",
    "\n",
    "Before inference, we inspect the RGB channels from the 4-band NAIP tensor (R, G, B, NIR). This confirms image orientation and content."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0825bb56",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(4, 4))\n",
    "plt.imshow(rgb)\n",
    "plt.axis('off')\n",
    "plt.title('NAIP RGB Tile')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b1935da",
   "metadata": {},
   "source": [
    "## Load Pre-trained Model\n",
    "\n",
    "We load `Unet_Weights.NAIP_RGBN_RESNET18_CHESAPEAKERSC` and initialize `unet` with the matching number of classes.\n",
    "\n",
    "Using the weights object ensures the model architecture and preprocessing transforms stay aligned with training configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "695e988b",
   "metadata": {},
   "outputs": [],
   "source": [
    "weights = Unet_Weights.NAIP_RGBN_RESNET18_CHESAPEAKERSC\n",
    "model = unet(weights=weights).eval()\n",
    "weights.meta"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11790f1c",
   "metadata": {},
   "source": [
    "## Run Inference\n",
    "\n",
    "We apply the weights-provided transforms, run the model in inference mode, convert logits to class probabilities, and threshold road probabilities at 0.5 for a binary mask."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3b773bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "preprocess = weights.transforms\n",
    "inputs = preprocess(image)\n",
    "with torch.inference_mode():\n",
    "    logits = model(inputs)\n",
    "pred = logits.softmax(dim=1)[0, 1]\n",
    "pred_mask = pred > 0.5\n",
    "target_mask = target > 0\n",
    "pred.shape, pred_mask.float().mean().item(), target_mask.float().mean().item()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d96b73b6",
   "metadata": {},
   "source": [
    "## Visualize Ground Truth vs Prediction\n",
    "\n",
    "We compare:\n",
    "\n",
    "* RGB input,\n",
    "* ground-truth road mask,\n",
    "* predicted road probability map,\n",
    "* thresholded predicted mask.\n",
    "\n",
    "This gives a quick qualitative check of segmentation behavior on a real sample tile."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1e5d4a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 4, figsize=(14, 4))\n",
    "ax[0].imshow(rgb)\n",
    "ax[0].set_title('RGB')\n",
    "ax[1].imshow(target_mask, cmap='gray')\n",
    "ax[1].set_title('Ground Truth Mask')\n",
    "ax[2].imshow(pred, cmap='magma')\n",
    "ax[2].set_title('Road Probability')\n",
    "ax[3].imshow(pred_mask, cmap='gray')\n",
    "ax[3].set_title('Predicted Mask > 0.5')\n",
    "for a in ax:\n",
    "    a.axis('off')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  }
 ],
 "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.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}