{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "76a9a18b",
   "metadata": {},
   "source": [
    "# Manual labeling for supervised classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f75d09d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import rioxarray as riox\n",
    "\n",
    "from pathlib import Path\n",
    "from polsarpro.io import open_netcdf_beam\n",
    "from polsarpro.util import pauli_rgb\n",
    "\n",
    "# optional import for progress bar\n",
    "from dask.diagnostics import ProgressBar\n",
    "\n",
    "\n",
    "# change to your data paths\n",
    "# original dataset\n",
    "input_alos_data = Path(\"/data/psp/test_files/SAN_FRANCISCO_ALOS1_slc.nc\")\n",
    "\n",
    "# output files \n",
    "output_test_dir = Path(\"/data/psp/\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f82ace0",
   "metadata": {},
   "source": [
    "## Open the image and make a Pauli RGB representation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7c0adce",
   "metadata": {},
   "outputs": [],
   "source": [
    "S = open_netcdf_beam(input_alos_data)\n",
    "rgb = pauli_rgb(S)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8a93514",
   "metadata": {},
   "source": [
    "## Convert to UINT8 and save as a Tiff file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d01b44de",
   "metadata": {},
   "outputs": [],
   "source": [
    "(255*rgb).astype(\"uint8\").rio.to_raster(output_test_dir / \"SF_Pauli.tif\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "747ed2ff",
   "metadata": {},
   "source": [
    "## Manual labeling\n",
    "Recommended approach:\n",
    "- Open the tif image in GIMP\n",
    "- Add one layer per label, use transparent background\n",
    "- Draw labels (any non-zero value) on each layer, avoid overlap bewteen different classes\n",
    "- Export each layer separately as a png file"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aefec81e",
   "metadata": {},
   "source": [
    "## Import hand-drawn labels\n",
    "\n",
    "- In the following we assume that labels are png files\n",
    "- We map the name of each layer to a class number\n",
    "- We produce an `xarray.DataArray` that may be read by the classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bea3a64",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.image import imread\n",
    "labels = {\"sea\": 1, \"vegetation\":2, \"urban\":3}\n",
    "\n",
    "# create a class map with integer type from the input data size\n",
    "class_map = np.zeros_like(S.hh.data.real, dtype=int)\n",
    "for name in labels:\n",
    "\n",
    "    # load class image\n",
    "    im_pth = Path(\"/data/psp/labels\") / f\"{name}.png\"\n",
    "    # if color image we need only one channel\n",
    "    im = imread(im_pth)[..., 0]\n",
    "    # add with corresponding label\n",
    "    class_map += labels[name] * (im > 0)\n",
    "\n",
    "# make a data array that can be read by the classifier\n",
    "da_class_map = xr.DataArray(class_map, dims=S.dims, coords=S.coords) \n",
    "\n",
    "# Optional: save array to a file using `to_netcdf`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "277f82fd",
   "metadata": {},
   "source": [
    "## Apply Wishart supervised classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2580db8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.diagnostics import ProgressBar\n",
    "from polsarpro.classification import wishart_supervised\n",
    "with ProgressBar():\n",
    "    cm = wishart_supervised(S, da_class_map).compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7cf03b3",
   "metadata": {},
   "source": [
    "## Visualize the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d2527f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import ListedColormap\n",
    "\n",
    "# Blue: sea, green: vegetation, orange: built-up\n",
    "colors = [\"#387EB9\", \"#4EAF49\", \"#FF7F00\"]\n",
    "colors2 = [\"#000000\", \"#387EB9\", \"#4EAF49\", \"#FF7F00\"]\n",
    "# Create the discrete colormap\n",
    "my_cmap = ListedColormap(colors)\n",
    "my_cmap2 = ListedColormap(colors2)\n",
    "\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.subplot(1,3,1)\n",
    "plt.imshow(cm.wishart_supervised_class[::8], interpolation=\"none\", cmap=my_cmap)\n",
    "plt.axis(\"off\")\n",
    "plt.title(\"Output labels\")\n",
    "plt.subplot(1,3,2)\n",
    "plt.imshow(class_map[::8], interpolation=\"none\", cmap=my_cmap2)\n",
    "plt.axis(\"off\")\n",
    "plt.title(\"Manual inputs\")\n",
    "plt.subplot(1,3,3)\n",
    "plt.imshow(rgb.transpose(\"y\",\"x\",\"band\")[::8], interpolation=\"bilinear\")\n",
    "plt.axis(\"off\")\n",
    "plt.title(\"Pauli RGB\")\n",
    "plt.tight_layout()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "psp",
   "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.12.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
