RiverMask_Cleaner.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "45fed24f",
   "metadata": {},
   "source": [
    "# River mask cleaner\n",
    "\n",
    "The following code takes DSWE-derived watermasks from a specified local path. It also finds the associated \"main_channel\" vector produced from the Calculate_CFI notebook. Then, it overlays the main channel vector with the channel mask and selects the wetted portions of the mask intersecting the main channel, thus eliminating detected surface water not connected to the river. The processed mask is placed in a \"Processed\" folder. If no \"Processed\" folder exists, a new folder with that name is generated. \n",
    "\n",
    "Author: James (Huck) Rees; PhD Student, UCSB Geography\n",
    "\n",
    "Date: July 12th, 2024"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83e5cb5b",
   "metadata": {},
   "source": [
    "## Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0a4cfbaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import rasterio\n",
    "from scipy.ndimage import label as ndimage_label\n",
    "import geopandas as gpd\n",
    "from shapely.geometry import box\n",
    "from shapely.geometry import Polygon, MultiPoint\n",
    "import os\n",
    "from rasterio.transform import from_origin\n",
    "import glob"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d31d5c0c",
   "metadata": {},
   "source": [
    "## Initialize functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "5e09ba76",
   "metadata": {},
   "outputs": [],
   "source": [
    "def import_raster(file_path):\n",
    "    \"\"\"\n",
    "    Imports a raster file from a specified local path.\n",
    "\n",
    "    Parameters:\n",
    "    file_path (str): The path to the raster file.\n",
    "\n",
    "    Returns:\n",
    "    dataset: A rasterio dataset object.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        dataset = rasterio.open(file_path)\n",
    "        return dataset\n",
    "    except rasterio.errors.RasterioIOError as e:\n",
    "        print(f\"Error opening raster file: {e}\")\n",
    "        return None\n",
    "\n",
    "def classify_wet_components(raster, min_pixels=200):\n",
    "    \"\"\"\n",
    "    Classify connected components of wet pixels in a 2-D water mask raster and eliminate small components.\n",
    "\n",
    "    Parameters:\n",
    "    raster: A rasterio dataset object with pixel values 0 (dry) and 1 (wet).\n",
    "    min_pixels (int): Minimum number of pixels for a component to be retained. Default is 200.\n",
    "\n",
    "    Returns:\n",
    "    classified_raster: A 2-D numpy array with unique integer values for each component.\n",
    "    \"\"\"\n",
    "    # Read the raster data into a numpy array\n",
    "    water_mask = raster.read(1)\n",
    "\n",
    "    # Label the connected components\n",
    "    structure = np.ones((3, 3), dtype=int)  # 8-connectivity\n",
    "    labeled_array, num_features = ndimage_label(water_mask, structure=structure)\n",
    "\n",
    "    # Eliminate components smaller than min_pixels\n",
    "    component_sizes = np.bincount(labeled_array.ravel())\n",
    "    for component_label, size in enumerate(component_sizes):\n",
    "        if size < min_pixels:\n",
    "            labeled_array[labeled_array == component_label] = 0\n",
    "\n",
    "    # Relabel the remaining components to have consecutive integer labels\n",
    "    labeled_array, num_features = ndimage_label(labeled_array > 0, structure=structure)\n",
    "\n",
    "    return labeled_array\n",
    "\n",
    "def import_and_reproject_shapefile(shapefile_path, raster_crs):\n",
    "    \"\"\"\n",
    "    Import a shapefile and reproject it to match the CRS of the raster.\n",
    "\n",
    "    Parameters:\n",
    "    shapefile_path (str): The path to the shapefile.\n",
    "    raster_crs (rasterio.crs.CRS): The CRS of the raster to match.\n",
    "\n",
    "    Returns:\n",
    "    gdf: A GeoDataFrame with the reprojected shapefile.\n",
    "    \"\"\"\n",
    "    # Import the shapefile\n",
    "    gdf = gpd.read_file(shapefile_path)\n",
    "\n",
    "    # Reproject the GeoDataFrame to match the raster CRS\n",
    "    gdf = gdf.to_crs(raster_crs)\n",
    "\n",
    "    return gdf\n",
    "\n",
    "def raster_to_gdf(labeled_array, transform):\n",
    "    \"\"\"\n",
    "    Convert a labeled raster array to a GeoDataFrame with polygons for each pixel.\n",
    "\n",
    "    Parameters:\n",
    "    labeled_array: A 2-D numpy array with labeled components.\n",
    "    transform: Affine transform for converting array coordinates to geographic coordinates.\n",
    "\n",
    "    Returns:\n",
    "    gdf: A GeoDataFrame with polygons for each pixel.\n",
    "    \"\"\"\n",
    "    rows, cols = labeled_array.shape\n",
    "    polygons = []\n",
    "    labels = []\n",
    "\n",
    "    for row in range(rows):\n",
    "        for col in range(cols):\n",
    "            value = labeled_array[row, col]\n",
    "            if value != 0:  # Skip background pixels\n",
    "                # Create a polygon for each pixel\n",
    "                minx, miny = transform * (col, row)\n",
    "                maxx, maxy = transform * (col + 1, row + 1)\n",
    "                polygons.append(box(minx, miny, maxx, maxy))\n",
    "                labels.append(value)\n",
    "\n",
    "    # Create a GeoDataFrame\n",
    "    gdf = gpd.GeoDataFrame({'geometry': polygons, 'label': labels})\n",
    "\n",
    "    return gdf\n",
    "\n",
    "def find_intersecting_pixels(raster_gdf, line_gdf):\n",
    "    \"\"\"\n",
    "    Find unique pixel values that a line intersects with.\n",
    "\n",
    "    Parameters:\n",
    "    raster_gdf: GeoDataFrame with raster pixels as polygons.\n",
    "    line_gdf: GeoDataFrame with polyline geometry.\n",
    "\n",
    "    Returns:\n",
    "    unique_labels: A set of unique pixel values that the line intersects.\n",
    "    \"\"\"\n",
    "    # Perform a spatial join to find intersecting pixels\n",
    "    intersections = gpd.sjoin(raster_gdf, line_gdf, how='inner', op='intersects')\n",
    "    unique_labels = set(intersections['label'])\n",
    "    \n",
    "    return unique_labels\n",
    "\n",
    "def filter_and_reclassify_raster(classified_raster, unique_labels, min_pixels=200):\n",
    "    \"\"\"\n",
    "    Filter the classified raster to include only specified components and reclassify pixels.\n",
    "    Eliminate components smaller than min_pixels.\n",
    "\n",
    "    Parameters:\n",
    "    classified_raster: A 2-D numpy array with labeled components.\n",
    "    unique_labels: A set of unique pixel values to retain as wet.\n",
    "    min_pixels (int): Minimum number of pixels for a component to be retained. Default is 200.\n",
    "\n",
    "    Returns:\n",
    "    final_raster: A 2-D numpy array with reclassified pixels.\n",
    "    \"\"\"\n",
    "    # Create a new array for the final raster\n",
    "    final_raster = np.zeros_like(classified_raster, dtype=np.int32)\n",
    "\n",
    "    # Set pixels that are in unique_labels to 1 (wet)\n",
    "    for unique_label in unique_labels:\n",
    "        final_raster[classified_raster == unique_label] = 1\n",
    "\n",
    "    # Label the connected components in the final raster\n",
    "    structure = np.ones((3, 3), dtype=int)  # 8-connectivity\n",
    "    labeled_array, num_features = ndimage_label(final_raster, structure=structure)\n",
    "\n",
    "    # Eliminate components smaller than min_pixels\n",
    "    component_sizes = np.bincount(labeled_array.ravel())\n",
    "    for component_label, size in enumerate(component_sizes):\n",
    "        if size < min_pixels:\n",
    "            final_raster[labeled_array == component_label] = 0\n",
    "\n",
    "    return final_raster\n",
    "\n",
    "def export_raster(raster, transform, crs, output_path):\n",
    "    \"\"\"\n",
    "    Export a raster array to a GeoTIFF file.\n",
    "\n",
    "    Parameters:\n",
    "    raster: A 2-D numpy array to export.\n",
    "    transform: Affine transform for the raster.\n",
    "    crs: Coordinate reference system of the raster.\n",
    "    output_path (str): The path to save the GeoTIFF file.\n",
    "    \"\"\"\n",
    "    with rasterio.open(\n",
    "        output_path,\n",
    "        'w',\n",
    "        driver='GTiff',\n",
    "        height=raster.shape[0],\n",
    "        width=raster.shape[1],\n",
    "        count=1,\n",
    "        dtype=raster.dtype,\n",
    "        crs=crs,\n",
    "        transform=transform,\n",
    "    ) as dst:\n",
    "        dst.write(raster, 1)\n",
    "\n",
    "def main(river_name, reach_range, year_range, base_raster_path, base_shapefile_path, min_pixels=200):\n",
    "    \"\"\"\n",
    "    Main function to process rasters and shapefile, and export the final rasters.\n",
    "    \"\"\"\n",
    "    river_folder = os.path.join(base_raster_path, river_name)\n",
    "    \n",
    "    if isinstance(reach_range, int):\n",
    "        reach_folders = [f\"reach_{reach_range}\"]\n",
    "    elif isinstance(reach_range, tuple):\n",
    "        reach_folders = [f\"reach_{i}\" for i in range(reach_range[0], reach_range[1] + 1)]\n",
    "    elif reach_range == \"All\":\n",
    "        reach_folders = [d for d in os.listdir(river_folder) if d.startswith(\"reach_\")]\n",
    "    else:\n",
    "        raise ValueError(\"Invalid reach_range input\")\n",
    "    \n",
    "    for reach_folder in reach_folders:\n",
    "        reach_number = int(reach_folder.split('_')[1])\n",
    "        reach_path = os.path.join(river_folder, reach_folder, \"Raw\")\n",
    "        \n",
    "        if isinstance(year_range, int):\n",
    "            years = [year_range]\n",
    "        elif isinstance(year_range, tuple):\n",
    "            years = range(year_range[0], year_range[1] + 1)\n",
    "        elif year_range == \"All\":\n",
    "            years = None  # Process all years\n",
    "        else:\n",
    "            raise ValueError(\"Invalid year_range input\")\n",
    "        \n",
    "        raster_files = glob.glob(os.path.join(reach_path, \"*.tif\"))\n",
    "        \n",
    "        for raster_path in raster_files:\n",
    "            raster_filename = os.path.basename(raster_path)\n",
    "            try:\n",
    "                year = int(raster_filename.split('_')[-4])\n",
    "                dswe_level = raster_filename.split('_')[-1][0]\n",
    "            except ValueError:\n",
    "                continue\n",
    "            \n",
    "            if years is not None and year not in years:\n",
    "                continue\n",
    "            \n",
    "            shapefile_path = os.path.join(base_shapefile_path, river_name, f\"reach_{reach_number}\", str(year), \"main_channel.shp\")\n",
    "            \n",
    "            raster = import_raster(raster_path)\n",
    "            classified_raster = classify_wet_components(raster, min_pixels=min_pixels)\n",
    "            main_channel = import_and_reproject_shapefile(shapefile_path, raster.crs)\n",
    "            raster_gdf = raster_to_gdf(classified_raster, raster.transform)\n",
    "            unique_labels = find_intersecting_pixels(raster_gdf, main_channel)\n",
    "            final_raster = filter_and_reclassify_raster(classified_raster, unique_labels, min_pixels=min_pixels)\n",
    "            \n",
    "            # Construct the output folder path\n",
    "            output_folder = os.path.join(base_raster_path, river_name, f\"reach_{reach_number}\", \"Processed\")\n",
    "            os.makedirs(output_folder, exist_ok=True)\n",
    "            \n",
    "            final_raster_filename = f\"{river_name}_reach_{reach_number}_{year}_DSWE_level_{dswe_level}_processed.tif\"\n",
    "            final_raster_path = os.path.join(output_folder, final_raster_filename)\n",
    "            \n",
    "            export_raster(final_raster, raster.transform, raster.crs, final_raster_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "02715a48",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\huckr\\AppData\\Local\\Temp\\ipykernel_18720\\245461602.py:10: FutureWarning: The `op` parameter is deprecated and will be removed in a future release. Please use the `predicate` parameter instead.\n",
      "  main(river_name, reach_range, year_range, base_raster_path, base_shapefile_path, min_pixels)\n",
      "C:\\Users\\huckr\\AppData\\Local\\Temp\\ipykernel_18720\\1756474094.py:108: UserWarning: CRS mismatch between the CRS of left geometries and the CRS of right geometries.\n",
      "Use `to_crs()` to reproject one of the input geometries to match the CRS of the other.\n",
      "\n",
      "Left CRS: None\n",
      "Right CRS: EPSG:4326\n",
      "\n",
      "  intersections = gpd.sjoin(raster_gdf, line_gdf, how='inner', op='intersects')\n",
      "C:\\Users\\huckr\\AppData\\Local\\Temp\\ipykernel_18720\\245461602.py:10: FutureWarning: The `op` parameter is deprecated and will be removed in a future release. Please use the `predicate` parameter instead.\n",
      "  main(river_name, reach_range, year_range, base_raster_path, base_shapefile_path, min_pixels)\n",
      "C:\\Users\\huckr\\AppData\\Local\\Temp\\ipykernel_18720\\1756474094.py:108: UserWarning: CRS mismatch between the CRS of left geometries and the CRS of right geometries.\n",
      "Use `to_crs()` to reproject one of the input geometries to match the CRS of the other.\n",
      "\n",
      "Left CRS: None\n",
      "Right CRS: EPSG:4326\n",
      "\n",
      "  intersections = gpd.sjoin(raster_gdf, line_gdf, how='inner', op='intersects')\n",
      "C:\\Users\\huckr\\AppData\\Local\\Temp\\ipykernel_18720\\245461602.py:10: FutureWarning: The `op` parameter is deprecated and will be removed in a future release. Please use the `predicate` parameter instead.\n",
      "  main(river_name, reach_range, year_range, base_raster_path, base_shapefile_path, min_pixels)\n",
      "C:\\Users\\huckr\\AppData\\Local\\Temp\\ipykernel_18720\\1756474094.py:108: UserWarning: CRS mismatch between the CRS of left geometries and the CRS of right geometries.\n",
      "Use `to_crs()` to reproject one of the input geometries to match the CRS of the other.\n",
      "\n",
      "Left CRS: None\n",
      "Right CRS: EPSG:4326\n",
      "\n",
      "  intersections = gpd.sjoin(raster_gdf, line_gdf, how='inner', op='intersects')\n",
      "C:\\Users\\huckr\\AppData\\Local\\Temp\\ipykernel_18720\\245461602.py:10: FutureWarning: The `op` parameter is deprecated and will be removed in a future release. Please use the `predicate` parameter instead.\n",
      "  main(river_name, reach_range, year_range, base_raster_path, base_shapefile_path, min_pixels)\n",
      "C:\\Users\\huckr\\AppData\\Local\\Temp\\ipykernel_18720\\1756474094.py:108: UserWarning: CRS mismatch between the CRS of left geometries and the CRS of right geometries.\n",
      "Use `to_crs()` to reproject one of the input geometries to match the CRS of the other.\n",
      "\n",
      "Left CRS: None\n",
      "Right CRS: EPSG:4326\n",
      "\n",
      "  intersections = gpd.sjoin(raster_gdf, line_gdf, how='inner', op='intersects')\n",
      "C:\\Users\\huckr\\AppData\\Local\\Temp\\ipykernel_18720\\245461602.py:10: FutureWarning: The `op` parameter is deprecated and will be removed in a future release. Please use the `predicate` parameter instead.\n",
      "  main(river_name, reach_range, year_range, base_raster_path, base_shapefile_path, min_pixels)\n",
      "C:\\Users\\huckr\\AppData\\Local\\Temp\\ipykernel_18720\\1756474094.py:108: UserWarning: CRS mismatch between the CRS of left geometries and the CRS of right geometries.\n",
      "Use `to_crs()` to reproject one of the input geometries to match the CRS of the other.\n",
      "\n",
      "Left CRS: None\n",
      "Right CRS: EPSG:4326\n",
      "\n",
      "  intersections = gpd.sjoin(raster_gdf, line_gdf, how='inner', op='intersects')\n",
      "C:\\Users\\huckr\\AppData\\Local\\Temp\\ipykernel_18720\\245461602.py:10: FutureWarning: The `op` parameter is deprecated and will be removed in a future release. Please use the `predicate` parameter instead.\n",
      "  main(river_name, reach_range, year_range, base_raster_path, base_shapefile_path, min_pixels)\n",
      "C:\\Users\\huckr\\AppData\\Local\\Temp\\ipykernel_18720\\1756474094.py:108: UserWarning: CRS mismatch between the CRS of left geometries and the CRS of right geometries.\n",
      "Use `to_crs()` to reproject one of the input geometries to match the CRS of the other.\n",
      "\n",
      "Left CRS: None\n",
      "Right CRS: EPSG:4326\n",
      "\n",
      "  intersections = gpd.sjoin(raster_gdf, line_gdf, how='inner', op='intersects')\n",
      "C:\\Users\\huckr\\AppData\\Local\\Temp\\ipykernel_18720\\245461602.py:10: FutureWarning: The `op` parameter is deprecated and will be removed in a future release. Please use the `predicate` parameter instead.\n",
      "  main(river_name, reach_range, year_range, base_raster_path, base_shapefile_path, min_pixels)\n",
      "C:\\Users\\huckr\\AppData\\Local\\Temp\\ipykernel_18720\\1756474094.py:108: UserWarning: CRS mismatch between the CRS of left geometries and the CRS of right geometries.\n",
      "Use `to_crs()` to reproject one of the input geometries to match the CRS of the other.\n",
      "\n",
      "Left CRS: None\n",
      "Right CRS: EPSG:4326\n",
      "\n",
      "  intersections = gpd.sjoin(raster_gdf, line_gdf, how='inner', op='intersects')\n",
      "C:\\Users\\huckr\\AppData\\Local\\Temp\\ipykernel_18720\\245461602.py:10: FutureWarning: The `op` parameter is deprecated and will be removed in a future release. Please use the `predicate` parameter instead.\n",
      "  main(river_name, reach_range, year_range, base_raster_path, base_shapefile_path, min_pixels)\n",
      "C:\\Users\\huckr\\AppData\\Local\\Temp\\ipykernel_18720\\1756474094.py:108: UserWarning: CRS mismatch between the CRS of left geometries and the CRS of right geometries.\n",
      "Use `to_crs()` to reproject one of the input geometries to match the CRS of the other.\n",
      "\n",
      "Left CRS: None\n",
      "Right CRS: EPSG:4326\n",
      "\n",
      "  intersections = gpd.sjoin(raster_gdf, line_gdf, how='inner', op='intersects')\n"
     ]
    }
   ],
   "source": [
    "# Input arguments\n",
    "river_name = 'Brahmaputra'  # e.g., 'Rakaia'\n",
    "reach_range = 25  # e.g., single int, tuple, or 'All'\n",
    "year_range = (1996, 2003)  # Can be a single year, a tuple, or 'All'\n",
    "base_raster_path = r\"C:\\Users\\huckr\\Desktop\\UCSB\\Dissertation\\Data\\RiverMapping\\RiverMasks\"\n",
    "base_shapefile_path = r\"C:\\Users\\huckr\\Desktop\\UCSB\\Dissertation\\Data\\RiverMapping\\Channels\"\n",
    "min_pixels = 200  # Minimum number of pixels for a component to be retained\n",
    "\n",
    "# Run the main function\n",
    "main(river_name, reach_range, year_range, base_raster_path, base_shapefile_path, min_pixels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b894e4c1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}