06_EMIT_Segment_Anything.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import leafmap\n",
    "import numpy as np\n",
    "import rioxarray as rxa\n",
    "import xarray as xr\n",
    "from dea_tools.spatial import xr_rasterize, xr_vectorize\n",
    "from fsspec.implementations.http import HTTPFileSystem\n",
    "from samgeo import SamGeo\n",
    "\n",
    "from emit_tools import emit_xarray\n",
    "from utils import gamma_adjust, get_earthdata_token, get_rgb_dataset, load_emit_granule"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Segment anything for EMIT data\n",
    "\n",
    "The [Segment Anything Model](https://segment-anything.com/) is an AI\n",
    "model to cut out features from an image.\n",
    "\n",
    "This notebook requires extra libraries installed and needs a GPU instance\n",
    "to work effectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "# Loading data can take around 3-4 minutes on a 100 Mbps connection\n",
    "\n",
    "# See README.md for instructions on how to get an Earthdata token\n",
    "token = get_earthdata_token()\n",
    "\n",
    "# Refer to the README.md for instructions on how to find different granule IDs\n",
    "granule = \"EMIT_L2A_RFL_001_20230316T045211_2307503_006\" # Canberra\n",
    "\n",
    "ds = load_emit_granule(granule, token)\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bands = {\n",
    "    \"red\": 650,\n",
    "    \"green\": 560,\n",
    "    \"blue\": 470,\n",
    "}\n",
    "\n",
    "dataset = xr.Dataset()\n",
    "\n",
    "for band, wavelength in bands.items():\n",
    "    data_array = xr.DataArray(\n",
    "        gamma_adjust(ds, wavelength, 0.4, replace_nans=False) * 255,\n",
    "        dims=('latitude', 'longitude'),\n",
    "        coords={'longitude': ds.longitude, 'latitude': ds.latitude}\n",
    "    ).astype(np.uint8)\n",
    "    dataset[band] = data_array\n",
    "\n",
    "# Rename coordinates\n",
    "dataset = dataset.rename({'longitude': 'x', 'latitude': 'y'})\n",
    "# see the structure\n",
    "print(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Python Segment Anything wants to work with files, which is inconvenient.\n",
    "image = \"scene_rgb.tif\"\n",
    "dataset.rio.to_raster(\"scene_rgb.tif\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show the saved image on a map\n",
    "m = leafmap.Map(center=[-42.9, 147.3], zoom=15)\n",
    "m.add_basemap(\"SATELLITE\")\n",
    "m.add_raster(image, layer_name=\"Image\")\n",
    "m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# When we're happy, delete the raw data to save memory\n",
    "del ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up our model\n",
    "sam_kwargs = {\n",
    "    \"points_per_side\": 32,\n",
    "    \"pred_iou_thresh\": 0.86,\n",
    "    \"stability_score_thresh\": 0.92,\n",
    "    \"crop_n_layers\": 1,\n",
    "    \"crop_n_points_downscale_factor\": 2,\n",
    "    \"min_mask_region_area\": 100,\n",
    "}\n",
    "\n",
    "sam = SamGeo(\n",
    "    model_type=\"vit_h\",\n",
    "    checkpoint=\"sam_vit_h_4b8939.pth\",\n",
    "    sam_kwargs=sam_kwargs,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sam.generate(image, \"scene_rgb_mask.tif\")\n",
    "sam.show_masks()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sam.show_masks()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Alex Test Environment",
   "language": "python",
   "name": "test"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}