SpaceNet SAR-RGB-Fusion.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SpaceNet Optical-SAR Fusion\n",
    "#### Created by Jason Brown and Jake Shermeyer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make the notebook interactive\n",
    "%matplotlib notebook\n",
    "\n",
    "# Required libraries\n",
    "import numpy as np\n",
    "from scipy import fftpack\n",
    "from skimage import io \n",
    "from skimage import transform\n",
    "import skimage.color as color\n",
    "from matplotlib import pyplot as plt \n",
    "from scipy.ndimage.filters import uniform_filter\n",
    "from scipy.ndimage.measurements import variance\n",
    "import spectral as spy\n",
    "import gdal\n",
    "import os\n",
    "import glob\n",
    "from tqdm.notebook import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The imagery data used here should be the corresponding SpaceNet 6 PS-RGB and SAR Intensity tiles. https://spacenet.ai/sn6-challenge/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Contrast Stretch function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def stretch(bands, lower_percent=1, higher_percent=98):\n",
    "    np.ma.array(bands, mask=np.isnan(bands))\n",
    "    out = np.zeros_like(bands)\n",
    "    a = 0 \n",
    "    b = 255 \n",
    "    c = np.percentile(bands, lower_percent)\n",
    "    d = np.percentile(bands, higher_percent)        \n",
    "    t = a + (bands - c) * (b - a) / (d - c)    \n",
    "    t[t<a] = a\n",
    "    t[t>b] = b\n",
    "    out = t\n",
    "    return out.astype(np.uint8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Lee Filter Function for Speckle Filtering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lee_filter(img, size):\n",
    "    img_mean = uniform_filter(img, (size, size))\n",
    "    img_sqr_mean = uniform_filter(img**2, (size, size))\n",
    "    img_variance = img_sqr_mean - img_mean**2\n",
    "    overall_variance = variance(img)\n",
    "    img_weights = img_variance / (img_variance + overall_variance)\n",
    "    img_output = img_mean + img_weights * (img - img_mean)\n",
    "    return img_output           "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PCT function for Polarimetric bands"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Principal components calculated from the Spectral package: https://github.com/spectralpython/spectral\n",
    "def pct_image(img, eigs):\n",
    "    pc = spy.principal_components(img)\n",
    "    ch = pc.reduce(eigs=eigs) # select the desired principal component\n",
    "    img_pc = ch.transform(img) # transform principal component into an image\n",
    "    return img_pc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build the Span image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def span_image(img):\n",
    "    img_sp = img[:,:,0]**2 + 2*abs(img[:,:,1]) + img[:,:,3]**2\n",
    "    return img_sp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fuse data sources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fusion(rgb, sar, method='hsv'):\n",
    "    # Get rgb_bands\n",
    "    R = rgb[:,:,0] # red\n",
    "    G = rgb[:,:,1] # green\n",
    "    B = rgb[:,:,2] # blue  \n",
    "    image = None\n",
    "    if method == 'simple_mean':\n",
    "        r = 0.5 * (0.5*R + sar)[:, :, np.newaxis]\n",
    "        g = 0.5 * (0.5*G + sar)[:, :, np.newaxis]\n",
    "        b = 0.5 * (0.5*B + sar)[:, :, np.newaxis]\n",
    "        image = np.concatenate([r,g,b], axis=2)   \n",
    "    if method == 'hsv':\n",
    "        hsv = color.rgb2hsv(rgb)\n",
    "        hsv[:,:,2] = sar \n",
    "        image = color.hsv2rgb(hsv)\n",
    "    return image    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Main processing code block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_multiband_geotiff(array, out_name, proj, geo, nodata=0, out_format=gdal.GDT_Byte, verbose=False):\n",
    "    \"\"\"Convert an array to an output georegistered geotiff.\n",
    "    Arguments\n",
    "    ---------\n",
    "    array : numpy array\n",
    "        A numpy array with a the shape: [Channels, X, Y] or [X, Y]\n",
    "    out_name : str\n",
    "        The output name and path for your image\n",
    "    proj : gdal projection\n",
    "        A projection, can be extracted from an image opened with gdal with image.GetProjection().  Can be set to None if no georeferencing is required.\n",
    "    geo : gdal geotransform\n",
    "        A gdal geotransform which indicates the position of the image on the earth in projection units. Can be set to None if no georeferencing is required.\n",
    "        Can be extracted from an image opened with gdal with image.GetGeoTransform()\n",
    "    nodata : int, default - 0\n",
    "        A value to set transparent for GIS systems. Can be set to None if the nodata value is not required.\n",
    "    out_format : gdalconst\n",
    "        https://gdal.org/python/osgeo.gdalconst-module.html\n",
    "        Must be one of the variables listed in the docs above\n",
    "    verbose : bool\n",
    "        A verbose output, printing all inputs and outputs to the function.  Useful for debugging.\n",
    "    \"\"\"\n",
    "    driver = gdal.GetDriverByName('GTiff')\n",
    "    if len(array.shape) == 2:\n",
    "        array = array[np.newaxis, ...]\n",
    "    os.makedirs(os.path.dirname(os.path.abspath(out_name)), exist_ok=True)\n",
    "    dataset = driver.Create(out_name, array.shape[2], array.shape[1], array.shape[0], out_format)\n",
    "    if verbose is True:\n",
    "        print(\"Array Shape, should be [Channels, X, Y] or [X,Y]:\", array.shape)\n",
    "        print(\"Output Name:\", out_name)\n",
    "        print(\"Projection:\", proj)\n",
    "        print(\"GeoTransform:\", geo)\n",
    "        print(\"NoData Value:\", nodata)\n",
    "        print(\"Bit Depth:\", out_format)\n",
    "    if proj is not None:\n",
    "        dataset.SetProjection(proj)\n",
    "    if geo is not None:\n",
    "        dataset.SetGeoTransform(geo)\n",
    "    if nodata is None:\n",
    "        for i, image in enumerate(array, 1):\n",
    "            dataset.GetRasterBand(i).WriteArray(image)\n",
    "        del dataset\n",
    "    else:\n",
    "        for i, image in enumerate(array, 1):\n",
    "            dataset.GetRasterBand(i).WriteArray(image)\n",
    "            dataset.GetRasterBand(i).SetNoDataValue(nodata)\n",
    "        del dataset\n",
    "\n",
    "def color_sar(sar_tiles_dir, rgb_tiles_dir, tiles_out_dir, search='.tif', method='hsv', span_or_pca='span', filtersize=3, nodata=0, add_4th=False):\n",
    "    \"\"\"Colorizes SAR data using co-located RGB imagery and outputs it to the tiles_out_dir.\n",
    "    Arguments\n",
    "    ---------\n",
    "    sar_tiles_dir : str\n",
    "        Path to the directory that contains SAR data\n",
    "    rgb_tiles_dir : str\n",
    "        Path to the directory that contains RGB data\n",
    "    tiles_out_dir : str\n",
    "        Path to the directory where colorized SAR will be output\n",
    "    search : str\n",
    "        File extension of imagery, defaults to '.tif'\n",
    "    method : str\n",
    "        Colorization method, should be 'hsv' or 'simple_mean'.  Defaults to 'hsv.'\n",
    "    span_or_pca : str\n",
    "        Calculate a span image or use pca to convert SAR to single channel.  Should be 'span' or 'pca.'\n",
    "        Defaults to 'span.'\n",
    "    nodata : int\n",
    "        A value to set transparent for GIS systems. Can be set to None if the nodata value is not required.\n",
    "        Defaults to 0.\n",
    "    filtersize : int\n",
    "        The filter size for a lee filter to reduce image noise.\n",
    "    add_4th : bool\n",
    "        Add an artificial 4th channel for working with networks that expect 4 channel inputs. Defaults to 'False'.\n",
    "    \"\"\"\n",
    "    if not os.path.exists(tiles_out_dir):\n",
    "        os.makedirs(tiles_out_dir)\n",
    "    os.chdir(sar_tiles_dir)\n",
    "    search2 = \"*\" + search\n",
    "    images = glob.glob(search2)\n",
    "    for image in tqdm(images):\n",
    "        sarimg = gdal.Open(image)\n",
    "        proj = sarimg.GetProjection()\n",
    "        geo = sarimg.GetGeoTransform()\n",
    "        sarimg = sarimg.ReadAsArray()\n",
    "        sarimg = np.swapaxes(sarimg,0,2)\n",
    "        sarimg = np.swapaxes(sarimg,0,1)\n",
    "        if sarimg.shape[2] > 1:\n",
    "            if span_or_pca == 'span':\n",
    "                span_img = span_image(sarimg)\n",
    "            elif span_or_pca == 'pca':\n",
    "                span_img = pct_image(sarimg, 0)\n",
    "            else:\n",
    "                print(\"Choose 'span' or 'pca' for span_or_pca\")\n",
    "        else:\n",
    "            span_img = sarimg[0,:,:]\n",
    "        new_image = image.split(\"SAR-Intensity\")[0] + \"PS-RGB\" + image.split(\"SAR-Intensity\")[1]\n",
    "        rgb_file = os.path.join(rgb_tiles_dir, new_image)\n",
    "        rgbimg = io.imread(rgb_file).astype(np.float32)\n",
    "        sarimg_hgm = transform.match_histograms(span_img, rgbimg[:,:,1])\n",
    "        lee_filt_img = lee_filter(sarimg_hgm, filtersize)\n",
    "        eo_sar_fusion = fusion(rgbimg, lee_filt_img, method=method) # see above for various methods available\n",
    "        eo_sar_fusion = np.swapaxes(eo_sar_fusion,1,0)\n",
    "        eo_sar_fusion = np.swapaxes(eo_sar_fusion,2,0)\n",
    "        output_image = image.split(\"SAR-Intensity\")[0] + \"SAR-Intensity-Colorized-HSV\" + image.split(\"SAR-Intensity\")[1]\n",
    "        output_image = os.path.join(tiles_out_dir, output_image)\n",
    "        if add_4th is True:\n",
    "            eo_sar_fusion = np.stack([eo_sar_fusion[0,:,:], eo_sar_fusion[1,:,:], eo_sar_fusion[2,:,:], eo_sar_fusion[0,:,:]])\n",
    "        create_multiband_geotiff(eo_sar_fusion, output_image, proj, geo, nodata=nodata)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Inputs and run function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sar_tiles_dir = \"../train/AOI_11_Rotterdam/SAR-Intensity\"\n",
    "rgb_tiles_dir = \"../train/AOI_11_Rotterdam/PS-RGB\"\n",
    "tiles_out_dir = \"../train/AOI_11_Rotterdam/Fused-RGB-SAR\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "color_sar(sar_tiles_dir, rgb_tiles_dir, tiles_out_dir, search='.tif', method='hsv', span_or_pca='span', filtersize=3, nodata=0, add_4th=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}