notebooks/crop_segmentation/02_visualize_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FarmVibes.AI Crop Segmentation - Data Exploration\n",
    "\n",
    "This notebook provides a self-contained workspace for visualizing and exploring the outputs of the dataset generation workflow for the crop segmentation task. \n",
    "\n",
    "We run a workflow that exposes the intermediate outputs of the dataset generation workflow, allowing us to visualize Sentinel-2 data, NDVI rasters and [Crop Data Layer](https://data.nal.usda.gov/dataset/cropscape-cropland-data-layer#:~:text=The%20Cropland%20Data%20Layer%20%28CDL%29%2C%20hosted%20on%20CropScape%2C,as%20well%20as%20boundary%2C%20water%20and%20road%20layers.) (CDL) maps provided by the FarmVibes.AI platform.\n",
    "\n",
    "\n",
    "### Micromamba environment setup\n",
    "Before running this notebook, let's build a micromamba environment. If you do not have micromamba installed, please follow the instructions from the [micromamba installation guide](https://mamba.readthedocs.io/en/latest/installation/micromamba-installation.html).\n",
    "\n",
    "```bash\n",
    "$ micromamba env create -f ./crop_env.yaml\n",
    "$ micromamba activate crop-seg\n",
    "```\n",
    "\n",
    "------------------------\n",
    "\n",
    "### Imports and utility methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Utility imports\n",
    "import json\n",
    "from datetime import datetime\n",
    "\n",
    "import geopandas as gpd\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import rasterio\n",
    "import yaml\n",
    "from matplotlib.colors import ListedColormap\n",
    "from notebook_lib.constants import CROP_INDICES\n",
    "from rasterio.mask import mask\n",
    "from rasterio.warp import reproject\n",
    "from shapely import wkt\n",
    "\n",
    "# FarmVibes.AI imports\n",
    "from vibe_core.client import get_default_vibe_client\n",
    "from vibe_notebook.plot import lw_plot, transparent_cmap\n",
    "from vibe_notebook.raster import read_raster, s2_to_img, spaceeye_to_img"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Retrieve the dataset with FarmVibes.AI platform\n",
    "\n",
    "We will retrieve the dataset from FarmVibes.AI cache by running a custom version of the `ml/dataset_generation/datagen_crop_segmentation` workflow that exposes the intermediate outputs (Sentinel-2 images, cloud and shadow mask, SpaceEye outputs) along with the NDVI rasters and CDL maps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = get_default_vibe_client()\n",
    "\n",
    "input_geometry_path = \"./input_region.wkt\"\n",
    "time_range = (datetime(2020, 1, 1), datetime(2020, 12, 31))\n",
    "\n",
    "# Reading the geometry file\n",
    "with open(input_geometry_path) as f:\n",
    "    geometry = wkt.load(f)\n",
    "\n",
    "# Reading the workflow\n",
    "with open(\"wf_viz_crop_seg.yaml\") as f:\n",
    "    wf_dict = yaml.safe_load(f)\n",
    "\n",
    "wf_run = client.run(\n",
    "    wf_dict, \"Crop Segmentation Workflow Visualization\", geometry=geometry, time_range=time_range\n",
    ")\n",
    "wf_run.monitor()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recovering the workflow outputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s2_rasters = wf_run.output[\"s2_rasters\"]\n",
    "s2_masks = wf_run.output[\"s2_masks\"]\n",
    "se_rasters = wf_run.output[\"se_rasters\"]\n",
    "ndvi_rasters = wf_run.output[\"ndvi\"]\n",
    "cdl_rasters = wf_run.output[\"cdl\"]\n",
    "\n",
    "print(f\"# of Sentinel-2 rasters: {len(s2_rasters)}\")\n",
    "print(f\"# of cloud and shadow masks: {len(s2_masks)}\")\n",
    "print(f\"# of SpaceEye rasters: {len(se_rasters)}\")\n",
    "print(f\"# of NDVI rasters: {len(ndvi_rasters)}\")\n",
    "print(f\"# of CDL maps: {len(cdl_rasters)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Exploration\n",
    "\n",
    "Let's select one Sentinel-2 raster and visualize it, along with the respective cloud and shadow mask and SpaceEye output.\n",
    "You can select other images by changing `select_idx` (between `0`  and `len(s2_rasters)`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "select_idx = 20\n",
    "select_idx = min(max(0, select_idx), len(s2_rasters) - 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s2 = s2_rasters[select_idx]\n",
    "\n",
    "\n",
    "s2_ar = read_raster(s2, geometry=geometry)[0]\n",
    "plt.figure(figsize=(8, 8))\n",
    "plt.title(f\"Sentinel-2: {s2.time_range[0].strftime('%Y-%m-%d')}\")\n",
    "plt.imshow(s2_to_img(s2_ar))\n",
    "plt.axis(\"off\")\n",
    "lw_plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Retrieving the cloud and shadow mask for the same date:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cloud_mask = [\n",
    "    m\n",
    "    for m in s2_masks\n",
    "    if m.time_range[0].day == s2.time_range[0].day\n",
    "    and m.time_range[0].month == s2.time_range[0].month\n",
    "][0]\n",
    "\n",
    "cloud_ar = read_raster(cloud_mask, geometry)[0]\n",
    "cloud_ar[cloud_ar == 100] = 0\n",
    "\n",
    "plt.figure(figsize=(8, 8))\n",
    "plt.imshow(s2_to_img(s2_ar))\n",
    "plt.imshow(cloud_ar[0], cmap=transparent_cmap(plt.cm.viridis))\n",
    "plt.title(f\"Mask: {cloud_mask.time_range[0].strftime('%Y-%m-%d')}\")\n",
    "plt.axis(\"off\")\n",
    "lw_plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And, finally, retrieve the SpaceEye output associated with them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spaceeye = [\n",
    "    se\n",
    "    for se in se_rasters\n",
    "    if se.time_range[0].day == s2.time_range[0].day\n",
    "    and se.time_range[0].month == s2.time_range[0].month\n",
    "][0]\n",
    "\n",
    "se_ar = read_raster(spaceeye, geometry)[0]\n",
    "\n",
    "plt.figure(figsize=(8, 8))\n",
    "plt.title(f\"SpaceEye: {spaceeye.time_range[0].strftime('%Y-%m-%d')}\")\n",
    "plt.imshow(spaceeye_to_img(se_ar))\n",
    "plt.axis(\"off\")\n",
    "lw_plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Side-by-side comparison:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(24, 8))\n",
    "\n",
    "plt.subplot(1, 3, 1)\n",
    "plt.imshow(s2_to_img(s2_ar))\n",
    "plt.title(\"Sentine-2 image\")\n",
    "plt.axis(\"off\")\n",
    "\n",
    "plt.subplot(1, 3, 2)\n",
    "plt.imshow(s2_to_img(s2_ar))\n",
    "plt.imshow(cloud_ar[0], cmap=transparent_cmap(plt.cm.viridis))\n",
    "plt.title(\"Overlaid cloud mask\")\n",
    "plt.axis(\"off\")\n",
    "\n",
    "plt.subplot(1, 3, 3)\n",
    "plt.imshow(spaceeye_to_img(se_ar))\n",
    "plt.title(\"SpaceEye image\")\n",
    "plt.axis(\"off\")\n",
    "\n",
    "lw_plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get a better sense of the differences between Sentinel-2 and SpaceEye images, let's view a few of them side by side:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_imgs_to_show = 5\n",
    "top_idx = min(select_idx + n_imgs_to_show, len(s2_rasters) - 1)\n",
    "\n",
    "for s2 in s2_rasters[select_idx:top_idx]:\n",
    "    plt.figure(figsize=(15, 8))\n",
    "    se_raster = sorted(se_rasters, key=lambda x: abs(s2.time_range[0] - x.time_range[0]))[0]\n",
    "    plt.subplot(1, 2, 1)\n",
    "    plt.title(f\"S-2: {s2.time_range[0].strftime('%Y-%m-%d')}\")\n",
    "    plt.imshow(s2_to_img(read_raster(s2, geometry)[0]))\n",
    "    plt.axis(\"off\")\n",
    "    plt.subplot(1, 2, 2)\n",
    "    plt.imshow(spaceeye_to_img(read_raster(se_raster)[0]))\n",
    "    plt.title(\"SpaceEye\")\n",
    "    plt.axis(\"off\")\n",
    "lw_plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The workflow computes the NDVI over the cloud-free image from SpaceEye:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ndvi = [\n",
    "    nd\n",
    "    for nd in ndvi_rasters\n",
    "    if nd.time_range[0].day == spaceeye.time_range[0].day\n",
    "    and nd.time_range[0].month == spaceeye.time_range[0].month\n",
    "][0]\n",
    "with rasterio.open(ndvi.raster_asset.url) as src:\n",
    "    ndvi_ar = src.read()\n",
    "\n",
    "\n",
    "plt.figure(figsize=(15, 8))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.imshow(spaceeye_to_img(se_ar))\n",
    "plt.title(\"SpaceEye\")\n",
    "plt.axis(\"off\")\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.imshow(ndvi_ar[0, :, :], vmin=0, vmax=1)\n",
    "plt.title(\"NDVI\")\n",
    "plt.colorbar()\n",
    "plt.axis(\"off\")\n",
    "lw_plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exploring the CDL maps\n",
    "Let's visualize the CDL map (with all crop indices colorized) and the binary (crop vs. non-crop) mask:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cdl = cdl_rasters[0]\n",
    "vis_asset = cdl.visualization_asset\n",
    "with open(vis_asset.local_path) as mtdt:\n",
    "    metadata = json.load(mtdt)\n",
    "    colormap = metadata[\"colormap\"]\n",
    "\n",
    "colormap = np.array(colormap) / 255.0\n",
    "colormap = ListedColormap(colormap.tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with rasterio.open(s2.raster_asset.url) as src:\n",
    "    proj_geom = gpd.GeoSeries(geometry, crs=\"epsg:4326\").to_crs(src.crs).iloc[0].envelope\n",
    "\n",
    "    ar, tr = mask(src, [proj_geom], crop=True)\n",
    "    cdl_ar = np.zeros_like(ar)\n",
    "    with rasterio.open(cdl.raster_asset.url) as src_w:\n",
    "        reproject(rasterio.band(src_w, 1), destination=cdl_ar, dst_transform=tr, dst_crs=src.crs)\n",
    "\n",
    "\n",
    "plt.figure(figsize=(15, 8))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.imshow(cdl_ar[0, :, :], cmap=colormap)\n",
    "plt.title(\"CDL Map\")\n",
    "plt.axis(\"off\")\n",
    "\n",
    "\n",
    "binary_cdl = np.isin(cdl_ar, CROP_INDICES)\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.imshow(binary_cdl[0, :, :])\n",
    "plt.title(\"Binary Crop Map\")\n",
    "plt.axis(\"off\")\n",
    "lw_plot()"
   ]
  }
 ],
 "metadata": {
  "description": "Visual exploration of the crop land segmentation dataset.",
  "disk_space": "",
  "kernelspec": {
   "display_name": "Python 3.8.13 ('crop-seg')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  },
  "name": "Crop land segmentation (2/4) - dataset visualization",
  "orig_nbformat": 4,
  "running_time": "",
  "tags": [
   "Crop Land Segmentation",
   "Working with Custom Workflows",
   "Remote Sensing"
  ],
  "vscode": {
   "interpreter": {
    "hash": "c4ed1382910929959896ae46210007cc63f041dd0c6059830101875a13ee5841"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}