07_S2_Segment_Anything.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "import pystac\n",
    "import rioxarray  # noqa F401\n",
    "import s3fs\n",
    "from odc.stac import configure_rio, load"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Segment anything for Sentinel-2 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.\n",
    "\n",
    "This example using Sentinel-2 data works reasonably well if we use 20 m\n",
    "pixels. 10 m pixels takes and exceedingly long time to run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_uri = \"s3://dea-public-data/baseline/ga_s2bm_ard_3/54/HYF/2023/01/25/20230125T014454/ga_s2bm_ard_3-2-1_54HYF_2023-01-25_final.stac-item.json\"\n",
    "\n",
    "# Region of interest to crop\n",
    "minx, maxx, miny, maxy = 143.5,144.0,-35.8,-35.3\n",
    "# Resolution\n",
    "RESOLUTION = 20\n",
    "\n",
    "# Open S3 object as a file using s3fs\n",
    "s3 = s3fs.S3FileSystem(anon=True)\n",
    "with s3.open(s3_uri, \"rt\") as f:\n",
    "    stac_dict = json.load(f)\n",
    "item = pystac.read_dict(stac_dict)\n",
    "\n",
    "# Load the data, telling rasterio to not sign requests\n",
    "configure_rio(\n",
    "    cloud_defaults=True,\n",
    "    aws={\"aws_unsigned\": True},\n",
    "    AWS_S3_ENDPOINT=\"s3.ap-southeast-2.amazonaws.com\",\n",
    ")\n",
    "data = load(\n",
    "    [item],\n",
    "    measurements=[\"green\", \"swir\", \"nir\"],\n",
    "    bbox=[minx, miny, maxx, maxy],\n",
    "    resolution=RESOLUTION,\n",
    "    stac_cfg={\n",
    "        \"ga_s2bm_ard_3\": {\n",
    "            \"aliases\": {\n",
    "                \"red\": \"nbart_red\",\n",
    "                \"green\": \"nbart_green\",\n",
    "                \"blue\": \"nbart_blue\",\n",
    "                \"nir\": \"nbart_nir_1\",\n",
    "                \"swir\": \"nbart_swir_2\",\n",
    "            }\n",
    "        }\n",
    "    },\n",
    "    chunks={}\n",
    ")\n",
    "\n",
    "# No need for time\n",
    "data = data.squeeze(\"time\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rescale red, green and blue bands from 0-10000 to 0-255\n",
    "\n",
    "# Copy data into a new xarray\n",
    "data_uint = data.copy(deep=True)\n",
    "\n",
    "data_uint[\"swir\"] = (data[\"swir\"] / 5000 * 255).astype(\"uint8\")\n",
    "data_uint[\"nir\"] = (data[\"nir\"] / 5000 * 255).astype(\"uint8\")\n",
    "data_uint[\"green\"] = (data[\"green\"] / 3000 * 255).astype(\"uint8\")\n",
    "\n",
    "image = \"s2_rgb.tif\"\n",
    "mask = \"s2_mask.tif\"\n",
    "data_uint.rio.to_raster(\"s2_rgb.tif\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from samgeo import SamGeo\n",
    "\n",
    "# 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",
    ")\n",
    "\n",
    "sam.generate(image, mask, foreground=False)\n",
    "sam.show_masks()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sam.save_masks(mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vector = \"segment.gpkg\"\n",
    "sam.tiff_to_gpkg(mask, vector, simplify_tolerance=None)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Alex Test Environment",
   "language": "python",
   "name": "test"
  },
  "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.6"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}