tutorial/05_Monitor_Downloads.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "707bb7bc-2991-4d25-8261-dede3fd111c3",
   "metadata": {},
   "source": [
    "# Monitor download state and task progress\n",
    "Run this notebook to visualise download state and monitor download tasks in progress on Google Earth Engine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d92f667-1d95-43e8-b47f-a928cdce5907",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Necessary imports\n",
    "import os\n",
    "os.environ['USE_PYGEOS'] = '0'\n",
    "import pandas as pd\n",
    "import geopandas as gpd\n",
    "import time\n",
    "import folium\n",
    "import geemap.foliumap as geemap\n",
    "import branca.colormap\n",
    "from tqdm.notebook import tqdm\n",
    "from datetime import datetime\n",
    "\n",
    "from db_utils import DB\n",
    "from dotenv import load_dotenv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d30dd68",
   "metadata": {},
   "source": [
    "## Load environment and project details\n",
    "\n",
    "As with the other notebooks, we load credentials and project details from a hidden ```.env``` file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e57b7c05",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load environment variables (including path to credentials) from '.env' file\n",
    "env_file_path = \"../.env\"\n",
    "\n",
    "# Uncomment for alternative version for Windows (r\"\" indicates raw string)\n",
    "#env_file_path = r\"C:/Users/User/floodmapper/.env\"\n",
    "\n",
    "assert load_dotenv(dotenv_path=env_file_path) == True, \"[ERR] Failed to load environment!\"\n",
    "assert \"GOOGLE_APPLICATION_CREDENTIALS\" in os.environ, \"[ERR] Missing $GOOGLE_APPLICATION_CREDENTIAL!\"\n",
    "assert \"GS_USER_PROJECT\" in os.environ, \"[ERR] Missing $GS_USER_PROJECT!\"\n",
    "key_file_path = os.environ[\"GOOGLE_APPLICATION_CREDENTIALS\"]\n",
    "assert os.path.exists(key_file_path), f\"[ERR] Google credential key file does not exist: \\n{key_file_path} \"\n",
    "assert \"ML4FLOODS_BASE_DIR\" in os.environ, \"[ERR] Missing $ML4FLOODS_BASE_DIR!\"\n",
    "base_path = os.environ[\"ML4FLOODS_BASE_DIR\"]\n",
    "assert os.path.exists(base_path), f\"[ERR] Base path does not exist: \\n{base_path} \"\n",
    "bucket_name = os.environ[\"BUCKET_URI\"]\n",
    "assert bucket_name is not None and bucket_name != \"\", f\"Bucket name not defined {bucket_name}\"\n",
    "print(\"[INFO] Successfully loaded FloodMapper environment.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21371a78",
   "metadata": {},
   "source": [
    "## Query the download state from the database\n",
    "\n",
    "**Set the name of the session here and run all remaining cells in order.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1aebf3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# EDIT THE NAME OF THE SESSION\n",
    "session_name = \"EMSR586\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b404539e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Connect to the database (point to the .env file for credentials)\n",
    "db_conn = DB(env_file_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf803cdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Query the geometry of the selected area\n",
    "query = (f\"SELECT DISTINCT sp.patch_name, ST_AsText(gr.geometry) \"\n",
    "         f\"FROM session_patches sp \"\n",
    "         f\"INNER JOIN world_grid gr \"\n",
    "         f\"ON sp.patch_name = gr.patch_name \"\n",
    "         f\"WHERE sp.session = %s ;\")\n",
    "data = (session_name,)\n",
    "grid_sel_df = db_conn.run_query(query, data, fetch=True)\n",
    "print(f\"[INFO] Returned {len(grid_sel_df)} rows.\")\n",
    "\n",
    "# Format the results into a correct GeoDataFrame\n",
    "grid_sel_df['geometry'] = gpd.GeoSeries.from_wkt(grid_sel_df['st_astext'])\n",
    "grid_sel_df.drop(['st_astext'], axis=1, inplace = True)\n",
    "grid_sel_gdf = gpd.GeoDataFrame(grid_sel_df, geometry='geometry', crs=\"EPSG:4326\")\n",
    "print(grid_sel_gdf.head(3))\n",
    "\n",
    "# Create an outline of the map area\n",
    "aoi_outline = grid_sel_gdf.geometry.unary_union\n",
    "aoi_outline_gdf = gpd.GeoDataFrame(geometry=[aoi_outline], crs=\"EPSG:4326\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91d1737e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Query the patches and their download status\n",
    "query = (f\"SELECT DISTINCT im.image_id, im.satellite, \"\n",
    "         f\"im.patch_name, im.status, ST_AsText(gr.geometry) \"\n",
    "         f\"FROM image_downloads im \"\n",
    "         f\"INNER JOIN world_grid gr \"\n",
    "         f\"ON im.patch_name = gr.patch_name \"\n",
    "         f\"INNER JOIN session_patches sp \"\n",
    "         f\"ON im.patch_name = sp.patch_name \"\n",
    "         f\"WHERE sp.session = %s ;\")\n",
    "data = (session_name,)\n",
    "grid_df = db_conn.run_query(query, data, fetch=True)\n",
    "print(f\"[INFO] Returned {len(grid_df)} rows.\")\n",
    "\n",
    "# Stop execution if no patches returned\n",
    "if len(grid_df) == 0:\n",
    "    print(f\"[INFO] No patches for session '{session_name}'.\")\n",
    "    raise KeyboardInterrupt\n",
    "\n",
    "# Format the results into a correct GeoDataFrame\n",
    "grid_df['geometry'] = gpd.GeoSeries.from_wkt(grid_df['st_astext'])\n",
    "grid_df.drop(['st_astext'], axis=1, inplace = True)\n",
    "grid_gdf = gpd.GeoDataFrame(grid_df, geometry='geometry', crs=\"EPSG:4326\")\n",
    "grid_gdf.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b754f2b",
   "metadata": {},
   "source": [
    "## Parse number of downloads and skipped files in each patch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23f569d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract the patch polygons\n",
    "geom = grid_gdf[[\"patch_name\", \"geometry\"]].drop_duplicates()\n",
    "geom = geom.set_index(\"patch_name\")\n",
    "\n",
    "# Count the downloaded and skipped files\n",
    "grid_dl_gdf = grid_gdf.loc[grid_gdf.status == 1]\n",
    "downloads = grid_dl_gdf.groupby(\"patch_name\").image_id.count()\n",
    "grid_skip_gdf = grid_gdf.loc[grid_gdf.status == 0]\n",
    "skipped = grid_skip_gdf.groupby(\"patch_name\").image_id.count()\n",
    "skipped = skipped.rename({\"image_id\": \"count\"})\n",
    "\n",
    "# Create a downloads gdf\n",
    "downloads_df = pd.concat([downloads, geom], axis = 1)\n",
    "downloads_df = downloads_df.rename(columns={\"image_id\": \"count\"})\n",
    "downloads_gdf = gpd.GeoDataFrame(downloads_df, geometry='geometry', crs=\"EPSG:4326\")\n",
    "downloads_gdf = downloads_gdf.fillna(0)\n",
    "\n",
    "# Create a skipped gdf\n",
    "skipped_df = pd.concat([skipped, geom], axis = 1)\n",
    "skipped_df = skipped_df.rename(columns={\"image_id\": \"count\"})\n",
    "skipped_gdf = gpd.GeoDataFrame(skipped_df, geometry='geometry', crs=\"EPSG:4326\")\n",
    "skipped_gdf = skipped_gdf.fillna(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa446f38",
   "metadata": {},
   "source": [
    "## Plot the number of downloaded files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65a7b63e",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# Define a style function to set the colours\n",
    "#cm = branca.colormap.LinearColormap(\n",
    "#    ['red', 'orange', 'yellow', 'cyan', 'blue', 'darkblue'],\n",
    "cm = branca.colormap.linear.YlOrRd_07.scale(\n",
    "    vmin=downloads_gdf[\"count\"].min(), \n",
    "    vmax=downloads_gdf[\"count\"].max())\n",
    "def style_fn(feature):\n",
    "    return {\n",
    "        'fillColor': cm(feature['properties']['count']),\n",
    "        'color': cm(feature['properties']['count']),\n",
    "        'weight': 0.5,\n",
    "        \"fillOpacity\": 0.5\n",
    "    }\n",
    "\n",
    "# Plot the patches colour-coded by number of downloads\n",
    "m = aoi_outline_gdf.explore(color=\"black\", style_kwds={\"fillOpacity\": 0.0, \"weight\": 3}, \n",
    "                            name=\"AoI Outline\", highlight=False)\n",
    "folium.GeoJson(downloads_gdf, \n",
    "               style_function=style_fn,\n",
    "               name=\"Downloaded Images\",\n",
    "               tooltip=folium.features.GeoJsonTooltip([\"count\"]),\n",
    "              ).add_to(m)\n",
    "\n",
    "# Add the colourmap, layer control and show\n",
    "m.add_child(cm)\n",
    "folium.LayerControl(collapsed=False).add_to(m)\n",
    "m"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "809c4889",
   "metadata": {},
   "source": [
    "## Plot the number of skipped files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d51a7fa3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a style function to set the colours\n",
    "#cm = branca.colormap.LinearColormap(\n",
    "#    ['red', 'yellow', 'green'],\n",
    "cm = branca.colormap.linear.YlOrRd_07.scale(\n",
    "    vmin=skipped_gdf[\"count\"].min(), \n",
    "    vmax=skipped_gdf[\"count\"].max())\n",
    "def style_fn(feature):\n",
    "    return {\n",
    "        'fillColor': cm(feature['properties']['count']),\n",
    "        'color': cm(feature['properties']['count']),\n",
    "        'weight': 0.5,\n",
    "        \"fillOpacity\": 0.5\n",
    "    }\n",
    "\n",
    "# Plot the patches colour-coded by number of downloads\n",
    "m = aoi_outline_gdf.explore(color=\"black\", style_kwds={\"fillOpacity\": 0.0, \"weight\": 3}, \n",
    "                            name=\"AoI Outline\", highlight=False)\n",
    "folium.GeoJson(skipped_gdf, \n",
    "               style_function=style_fn,\n",
    "               name=\"Skipped Images\",\n",
    "               tooltip=folium.features.GeoJsonTooltip([\"count\"]),\n",
    "              ).add_to(m)\n",
    "\n",
    "# Add the colourmap, layer control and show\n",
    "m.add_child(cm)\n",
    "folium.LayerControl(collapsed=False).add_to(m)\n",
    "m"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aedaab3f",
   "metadata": {},
   "source": [
    "## Monitor the task count and progress\n",
    "\n",
    "The cell below can be run to check the number of tasks being tracked by the database. Note that the ```01_download_images.py``` script must remain running for this notebook to work correctly.\n",
    "\n",
    "Individual tasks can also be viewed on this GEE [task tracking page](https://code.earthengine.google.com/tasks), provided you are logged Google Earth Engine via your web browser. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "944750ad-bab7-422e-9714-e49522d2e029",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to query the task status in the database\n",
    "def query_tasks():\n",
    "    query = (f\"SELECT description, state_code \"\n",
    "             f\"FROM gee_task_tracker \"\n",
    "             f\"WHERE session = %s;\")\n",
    "    data = (session_name,)\n",
    "    tasks_df = db_conn.run_query(query, data, fetch=True)\n",
    "    tasks_grp = tasks_df.groupby(\"state_code\")\n",
    "    \n",
    "    # Get the total number of tasks\n",
    "    num_tasks = len(tasks_df)\n",
    "    \n",
    "    # Get the number completed\n",
    "    num_completed = 0\n",
    "    if \"COMPLETED\" in tasks_grp.groups.keys():\n",
    "        num_completed = tasks_grp.count().loc[\"COMPLETED\"][0]\n",
    "\n",
    "    return num_tasks, num_completed\n",
    "\n",
    "# Query the total number of tasks being tracked\n",
    "num_tasks_old, _ = query_tasks()\n",
    "print(f\"[INFO] There are currently {num_tasks_old} tasks being tracked.\")\n",
    "\n",
    "# Initialise a progress bar\n",
    "pbar = tqdm(total=num_tasks_old,\n",
    "           dynamic_ncols=True,\n",
    "           leave=False,\n",
    "           position=0,\n",
    "           desc=\"Task Progress\")\n",
    "\n",
    "# Query the status every <interval_s> and update bar\n",
    "interval_s = 1\n",
    "num_completed_old = 0\n",
    "while True:\n",
    "    num_tasks, num_completed = query_tasks()\n",
    "    if num_tasks > num_tasks_old:\n",
    "        print(f\"[WARN] Total number of tasks has changed!\\n\"\n",
    "              f\"       Restart cell after all tasks have been submitted.\")\n",
    "        break\n",
    "    num_increment = num_completed - num_completed_old\n",
    "    num_completed_old = num_completed\n",
    "    pbar.update(num_increment)\n",
    "    now = datetime.now()\n",
    "    poll_time = now.strftime(\"%H:%M:%S\")\n",
    "    print(f\"[INFO] Last polling time: {poll_time}\", end=\"\\r\")\n",
    "    \n",
    "    # Exit if all tasks completed\n",
    "    if num_tasks == 0 or num_completed == num_tasks:\n",
    "        print(f\"[INFO] No active tasks remaining.\")\n",
    "        break\n",
    "    \n",
    "    # Increment after interval\n",
    "    time.sleep(interval_s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83fe0a74",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecbb9c06-2b6a-4e37-899e-3a77a2a2cae5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (floodmapper)",
   "language": "python",
   "name": "floodmapper"
  },
  "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}