#!/usr/bin/env python
# coding: utf-8

# # Tree crown delineation using detectreeRGB
# 
# {bdg-primary}`Forest`
# {bdg-secondary}`Modelling`
# {bdg-warning}`Standard`
# {bdg-info}`Python`
# 
# <p align="left">
#     <a href="https://github.com/eds-book-gallery/94486a7f-e046-461f-bbb9-334ec7b57040/blob/main/LICENSE">
#         <img alt="license" src="https://img.shields.io/badge/license-MIT-yellow.svg">
#     </a>
#     <a href="https://github.com/eds-book-gallery/94486a7f-e046-461f-bbb9-334ec7b57040/actions/workflows/render.yaml">
#         <img alt="render" src="https://github.com/eds-book-gallery/94486a7f-e046-461f-bbb9-334ec7b57040/actions/workflows/render.yaml/badge.svg">
#     </a>
#     <a href="https://github.com/alan-turing-institute/environmental-ds-book/pull/4">
#         <img alt="review" src="https://img.shields.io/badge/view-review-purple">
#     </a>
#     <br/>
# </p>
# 
# <p align="left">
#     <a href="http://mybinder.org/v2/gh/eds-book-gallery/94486a7f-e046-461f-bbb9-334ec7b57040/main?labpath=notebook.ipynb">
#         <img alt="binder" src="https://mybinder.org/badge_logo.svg">
#     </a>
#     <a href="https://replay.notebooks.egi.eu/v2/gh/eds-book-gallery/94486a7f-e046-461f-bbb9-334ec7b57040/main?labpath=notebook.ipynb">
#         <img alt="binder" src="https://img.shields.io/badge/launch-EGI%20binder-F5A252.svg?logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAFkAAABZCAMAAABi1XidAAAB8lBMVEX///9XmsrmZYH1olJXmsr1olJXmsrmZYH1olJXmsr1olJXmsrmZYH1olL1olJXmsr1olJXmsrmZYH1olL1olJXmsrmZYH1olJXmsr1olL1olJXmsrmZYH1olL1olJXmsrmZYH1olL1olL0nFf1olJXmsrmZYH1olJXmsq8dZb1olJXmsrmZYH1olJXmspXmspXmsr1olL1olJXmsrmZYH1olJXmsr1olL1olJXmsrmZYH1olL1olLeaIVXmsrmZYH1olL1olL1olJXmsrmZYH1olLna31Xmsr1olJXmsr1olJXmsrmZYH1olLqoVr1olJXmsr1olJXmsrmZYH1olL1olKkfaPobXvviGabgadXmsqThKuofKHmZ4Dobnr1olJXmsr1olJXmspXmsr1olJXmsrfZ4TuhWn1olL1olJXmsqBi7X1olJXmspZmslbmMhbmsdemsVfl8ZgmsNim8Jpk8F0m7R4m7F5nLB6jbh7jbiDirOEibOGnKaMhq+PnaCVg6qWg6qegKaff6WhnpKofKGtnomxeZy3noG6dZi+n3vCcpPDcpPGn3bLb4/Mb47UbIrVa4rYoGjdaIbeaIXhoWHmZYHobXvpcHjqdHXreHLroVrsfG/uhGnuh2bwj2Hxk17yl1vzmljzm1j0nlX1olL3AJXWAAAAbXRSTlMAEBAQHx8gICAuLjAwMDw9PUBAQEpQUFBXV1hgYGBkcHBwcXl8gICAgoiIkJCQlJicnJ2goKCmqK+wsLC4usDAwMjP0NDQ1NbW3Nzg4ODi5+3v8PDw8/T09PX29vb39/f5+fr7+/z8/Pz9/v7+zczCxgAABC5JREFUeAHN1ul3k0UUBvCb1CTVpmpaitAGSLSpSuKCLWpbTKNJFGlcSMAFF63iUmRccNG6gLbuxkXU66JAUef/9LSpmXnyLr3T5AO/rzl5zj137p136BISy44fKJXuGN/d19PUfYeO67Znqtf2KH33Id1psXoFdW30sPZ1sMvs2D060AHqws4FHeJojLZqnw53cmfvg+XR8mC0OEjuxrXEkX5ydeVJLVIlV0e10PXk5k7dYeHu7Cj1j+49uKg7uLU61tGLw1lq27ugQYlclHC4bgv7VQ+TAyj5Zc/UjsPvs1sd5cWryWObtvWT2EPa4rtnWW3JkpjggEpbOsPr7F7EyNewtpBIslA7p43HCsnwooXTEc3UmPmCNn5lrqTJxy6nRmcavGZVt/3Da2pD5NHvsOHJCrdc1G2r3DITpU7yic7w/7Rxnjc0kt5GC4djiv2Sz3Fb2iEZg41/ddsFDoyuYrIkmFehz0HR2thPgQqMyQYb2OtB0WxsZ3BeG3+wpRb1vzl2UYBog8FfGhttFKjtAclnZYrRo9ryG9uG/FZQU4AEg8ZE9LjGMzTmqKXPLnlWVnIlQQTvxJf8ip7VgjZjyVPrjw1te5otM7RmP7xm+sK2Gv9I8Gi++BRbEkR9EBw8zRUcKxwp73xkaLiqQb+kGduJTNHG72zcW9LoJgqQxpP3/Tj//c3yB0tqzaml05/+orHLksVO+95kX7/7qgJvnjlrfr2Ggsyx0eoy9uPzN5SPd86aXggOsEKW2Prz7du3VID3/tzs/sSRs2w7ovVHKtjrX2pd7ZMlTxAYfBAL9jiDwfLkq55Tm7ifhMlTGPyCAs7RFRhn47JnlcB9RM5T97ASuZXIcVNuUDIndpDbdsfrqsOppeXl5Y+XVKdjFCTh+zGaVuj0d9zy05PPK3QzBamxdwtTCrzyg/2Rvf2EstUjordGwa/kx9mSJLr8mLLtCW8HHGJc2R5hS219IiF6PnTusOqcMl57gm0Z8kanKMAQg0qSyuZfn7zItsbGyO9QlnxY0eCuD1XL2ys/MsrQhltE7Ug0uFOzufJFE2PxBo/YAx8XPPdDwWN0MrDRYIZF0mSMKCNHgaIVFoBbNoLJ7tEQDKxGF0kcLQimojCZopv0OkNOyWCCg9XMVAi7ARJzQdM2QUh0gmBozjc3Skg6dSBRqDGYSUOu66Zg+I2fNZs/M3/f/Grl/XnyF1Gw3VKCez0PN5IUfFLqvgUN4C0qNqYs5YhPL+aVZYDE4IpUk57oSFnJm4FyCqqOE0jhY2SMyLFoo56zyo6becOS5UVDdj7Vih0zp+tcMhwRpBeLyqtIjlJKAIZSbI8SGSF3k0pA3mR5tHuwPFoa7N7reoq2bqCsAk1HqCu5uvI1n6JuRXI+S1Mco54YmYTwcn6Aeic+kssXi8XpXC4V3t7/ADuTNKaQJdScAAAAAElFTkSuQmCC">
#     </a>
#     <br/>
# </p>
# 
# <p align="left">
#     <a href="https://w3id.org/ro-id/94486a7f-e046-461f-bbb9-334ec7b57040">
#         <img alt="RoHub" src="https://img.shields.io/badge/RoHub-FAIR_Executable_Research_Object-2ea44f?logo=Open+Access&logoColor=blue">
#     </a>
#     <a href="https://zenodo.org/badge/latestdoi/493595755">
#         <img alt="doi" src="https://zenodo.org/badge/493595755.svg">
#     </a>
# </p>
# 

# ## Context
# ### Purpose
# Accurately delineating trees using `detectron2`, a library that provides state-of-the-art deep learning detection and segmentation algorithms.
# 
# ### Modelling approach
# An established deep learning model, **Mask R-CNN** was deployed from `detectron2` library to delineate tree crowns accurately. A pre-trained model, named `detectreeRGB`, is provided to predict the location and extent of tree crowns from a top-down RGB image, captured by drone, aircraft or satellite. `detectreeRGB` was implemented in `python` 3.8 using `pytorch` v1.7.1 and `detectron2` v0.5. Further details can be found in the [repository documentation](https://github.com/shmh40/detectreeRGB).
# 
# ### Highlights
# * [detectreeRGB](https://github.com/shmh40/detectreeRGB) advances the state-of-the-art in tree identification from RGB images by delineating exactly the extent of the tree crown.
# * We demonstrate how to apply the pretrained model to a sample image fetched from a Zenodo repository.
# * Our pre-trained model was developed using aircraft images of tropical forests in Malaysia.
# * The model can be further trained using the user's own images.
# 
# ### Contributions
# 
# #### Notebook
# * Sebastian H. M. Hickman (author), University of Cambridge, [@shmh40](https://github.com/shmh40)
# * Alejandro Coca-Castro (reviewer), The Alan Turing Institute, [@acocac](https://github.com/acocac)
# 
# #### Modelling codebase
# * Sebastian H. M. Hickman (author), University of Cambridge [@shmh40](https://github.com/shmh40)
# * James G. C. Ball (contributor), University of Cambridge [@PatBall1](https://github.com/PatBall1)
# * David A. Coomes (contributor), University of Cambridge
# * Toby Jackson (contributor), University of Cambridge
# 
# #### Modelling funding
# The project was supported by the UKRI Centre for Doctoral Training in Application of Artificial Intelligence to the study of Environmental Risks ([AI4ER](https://ai4er-cdt.esc.cam.ac.uk/)) (EP/S022961/1).
# 
# :::{note}
# The authors acknowledge the authors of the [Detectron2](https://github.com/facebookresearch/detectron2) package which provides the Mask R-CNN architecture.
# :::

# ## Install and load libraries

# In[ ]:


get_ipython().system('pip -q install git+https://github.com/facebookresearch/detectron2.git')


# In[ ]:


import cv2
from PIL import Image
import os
import numpy as np
import urllib.request
import glob

# intake library and plugin
import intake

# geospatial libraries
import geopandas as gpd

from rasterio.transform import from_origin
import rasterio.features

import fiona

from shapely.geometry import shape, mapping, box
from shapely.geometry.multipolygon import MultiPolygon

# machine learning libraries
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.utils.visualizer import Visualizer, ColorMode
from detectron2.config import get_cfg
from detectron2.engine import DefaultTrainer

# visualisation
import holoviews as hv
import geoviews.tile_sources as gts
import matplotlib.pyplot as plt

import hvplot.pandas
import hvplot.xarray

import pooch

import warnings
warnings.filterwarnings(action='ignore')

hv.extension('bokeh', width=100)


# ## Set folder structure

# In[ ]:


# Define the project main folder
notebook_folder = './notebook'

# Set the folder structure
config = {
    'in_geotiff': os.path.join(notebook_folder, 'input','tiff'),
    'in_png': os.path.join(notebook_folder, 'input','png'),
    'model': os.path.join(notebook_folder, 'model'),
    'out_geotiff': os.path.join(notebook_folder, 'output','raster'),
    'out_shapefile': os.path.join(notebook_folder, 'output','vector'),
}

# List comprehension for the folder structure code
[os.makedirs(val) for key, val in config.items() if not os.path.exists(val)]


# ## Load and prepare input image

# ### Fetch a GeoTIFF file of aerial forest imagery using `pooch`
# 
# Let's fetch a sample aerial image from a [Zenodo repository](https://zenodo.org/record/5494629#.YWQQetnMKjA).

# In[ ]:


pooch.retrieve(
    url="doi:10.5281/zenodo.5494629/Sep_2014_RGB_602500_646600.tif",
    known_hash="md5:77a3b57f5f5946504ec520d1e793f250",
    path=notebook_folder,
    fname="input/tiff/Sep_2014_RGB_602500_646600.tif"
)


# In[ ]:


# set catalogue location
catalog_file = os.path.join(notebook_folder, 'catalog.yaml')

with open(catalog_file, 'w') as f:
    f.write('''
sources:
  sepilok_rgb:
    driver: rasterio
    description: 'NERC RGB images of Sepilok, Sabah, Malaysia (collection)'
    args:
      urlpath: "notebook/input/tiff/Sep_2014_RGB_602500_646600.tif"
      ''')


# In[ ]:


cat_tc = intake.open_catalog(catalog_file)


# In[ ]:


tc_rgb = cat_tc["sepilok_rgb"].to_dask()


# ### Inspect the aerial image
# 
# Let's investigate the `data-array`, what is the shape? Bounds? Bands? CRS?

# In[ ]:


print('shape =', tc_rgb.shape,',', 'and number of bands =', tc_rgb.count)


# ### Save the RGB bands of the GeoTIFF file as a PNG
# 
# **Mask R-CNN** requires images in `png` format. Let's export the RGB bands to a `png` file.

# In[ ]:


minx = 602500
miny = 646600

R = tc_rgb[0]
G = tc_rgb[1]
B = tc_rgb[2]
    
# stack up the bands in an order appropriate for saving with cv2, then rescale to the correct 0-255 range for cv2

# you will have to change the rescaling depending on the values of your tiff!
rgb = np.dstack((R,G,B)) # BGR for cv2
rgb_rescaled = 255*rgb/65535 # scale to image
    
# save this as png, named with the origin of the specific tile - change the filepath!
filepath = config['in_png'] + '/' + 'tile_' + str(minx) + '_' + str(miny) + '.png'
cv2.imwrite(filepath, rgb_rescaled)


# ### Read in and display the PNG file

# In[ ]:


im = cv2.imread(filepath)
plot_input = plt.figure(figsize=(15,15))
plt.imshow(Image.fromarray(im))
plt.title('Input image',fontsize='xx-large')
plt.axis('off')
plt.show()


# ## Modelling
# 
# ### Download the pretrained model

# In[ ]:


# define the URL to retrieve the model
fn = 'model_final.pth'
url = f'https://zenodo.org/record/5515408/files/{fn}?download=1'

urllib.request.urlretrieve(url, config['model'] + '/' + fn)


# ### Settings of `detectron2` config
# 
# The following lines allow configuring the main settings for predictions and load them into a `DefaultPredictor` object.

# In[ ]:


cfg = get_cfg()

# if you want to make predictions using a CPU, run the following line. If using GPU, hash it out.
cfg.MODEL.DEVICE='cpu'

# model and hyperparameter selection
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml"))
cfg.DATALOADER.NUM_WORKERS = 2
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1

### path to the saved pre-trained model weights
cfg.MODEL.WEIGHTS = config['model'] + '/model_final.pth'

# set confidence threshold at which we predict
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.15

#### Settings for predictions using detectron config

predictor = DefaultPredictor(cfg)


# ## Outputs
# 
# ### Showing the predictions from `detectreeRGB`

# In[ ]:


outputs = predictor(im)
v = Visualizer(im[:, :, ::-1], scale=1.5, instance_mode=ColorMode.IMAGE_BW)   # remove the colors of unsegmented pixels
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
image = cv2.cvtColor(v.get_image()[:, :, :], cv2.COLOR_BGR2RGB)

plot_predictions = plt.figure(figsize=(15,15))
plt.imshow(Image.fromarray(image))
plt.title('Predictions',fontsize='xx-large')
plt.axis('off')
plt.show()


# ### Convert predictions into geospatial files
# 
# #### To GeoTIFF

# In[ ]:


mask_array = outputs['instances'].pred_masks.cpu().numpy()

# get confidence scores too 
mask_array_scores = outputs['instances'].scores.cpu().numpy()

num_instances = mask_array.shape[0]
mask_array_instance = []
output = np.zeros_like(mask_array) 

mask_array_instance.append(mask_array)
output = np.where(mask_array_instance[0] == True, 255, output)
fresh_output = output.astype(float)
x_scaling = 140/fresh_output.shape[1]
y_scaling = 140/fresh_output.shape[2]
# this is an affine transform. This needs to be altered significantly.
transform = from_origin(int(filepath[-17:-11])-20, int(filepath[-10:-4])+120, y_scaling, x_scaling)

output_raster = config['out_geotiff'] + '/' + 'predicted_rasters_' + filepath[-17:-4]+ '.tif'

new_dataset = rasterio.open(output_raster, 'w', driver='GTiff',
                                height = fresh_output.shape[1], width = fresh_output.shape[2], count = fresh_output.shape[0],
                                dtype=str(fresh_output.dtype),
                                crs='+proj=utm +zone=50 +datum=WGS84 +units=m +no_defs',  
                                transform=transform)

new_dataset.write(fresh_output)
new_dataset.close()


# #### To shapefile

# In[ ]:


# Read input band with Rasterio
    
with rasterio.open(output_raster) as src:
    shp_schema = {'geometry': 'MultiPolygon','properties': {'pixelvalue': 'int', 'score': 'float'}}    

    crs = src.crs
    for i in range(src.count):
        src_band = src.read(i+1)
        src_band = np.float32(src_band)
        conf = mask_array_scores[i]
        # Keep track of unique pixel values in the input band
        unique_values = np.unique(src_band)
        # Polygonize with Rasterio. `shapes()` returns an iterable
        # of (geom, value) as tuples
        shapes = list(rasterio.features.shapes(src_band, transform=src.transform))

        if i == 0:
            with fiona.open(config['out_shapefile'] + '/predicted_polygons_' + filepath[-17:-4] + '_' + str(0) + '.shp', 'w', 'ESRI Shapefile',
                            shp_schema) as shp:
                polygons = [shape(geom) for geom, value in shapes if value == 255.0]                                        
                multipolygon = MultiPolygon(polygons)
                        # simplify not needed here
                        #multipolygon = multipolygon_a.simplify(0.1, preserve_topology=False)                    
                shp.write({
                          'geometry': mapping(multipolygon),
                          'properties': {'pixelvalue': int(unique_values[1]), 'score': float(conf)} 
                           })
        else:
            with fiona.open(config['out_shapefile'] + '/predicted_polygons_' + filepath[-17:-4] + '_' + str(0)+'.shp', 'a', 'ESRI Shapefile',
                            shp_schema) as shp:
                polygons = [shape(geom) for geom, value in shapes if value == 255.0]                                        
                multipolygon = MultiPolygon(polygons)
                        # simplify not needed here
                        #multipolygon = multipolygon_a.simplify(0.1, preserve_topology=False)                    
                shp.write({
                          'geometry': mapping(multipolygon),
                          'properties': {'pixelvalue': int(unique_values[1]), 'score': float(conf)} 
                           })


# ## Interactive map to visualise the exported files
# 
# ### Plot tree crown delineation predictions and scores

# In[ ]:


# load and plot polygons
in_shp = glob.glob(config['out_shapefile'] + '/*.shp')

poly_df = gpd.read_file(in_shp[0])

plot_vector = poly_df.hvplot(hover_cols=['score'], legend=False).opts(fill_color=None,line_color=None,alpha=0.5, width=800, height=600, xaxis=None, yaxis=None)
plot_vector


# ### Plot the exported GeoTIFF file

# In[ ]:


# load and plot RGB image
r = tc_rgb.sel(band=[1,2,3])

normalized = r/(r.quantile(.99,skipna=True)/255)

mask = normalized.where(normalized < 255)

int_arr = mask.astype(int)

plot_rgb = int_arr.astype('uint8').hvplot.rgb(
    x='x', y='y', bands='band', data_aspect=0.8, hover=False, legend=False, rasterize=True, xaxis=None, yaxis=None, title='Tree crown delineation by detectreeRGB'
)


# ### Overlay the prediction labels and image
# 
# Note we have some artifacts in the RGB image due to the transformations using the normalization procedure.

# In[ ]:


plot_predictions_interactive = plot_rgb * plot_vector
plot_predictions_interactive


# ### Save plot

# In[ ]:


hvplot.save(plot_predictions_interactive, notebook_folder + '/interactive_predictions.html')


# ## Summary
# 
# We have read in a raster, chosen a tile and made predictions on it. These predictions can then be transformed to shapefiles and examined in GIS software!
# 
# * We made the predictions on the `png` using a pretrained **Mask R-CNN** model, `detectreeRGB`.
# * The outputs showed the model capability to detect and delineate tree crowns from aerial imagery.
# * We then extracted our predictions, added the geospatial location back in, and exported them as shapefiles including the confidence score assigned to each prediction by the model.
# * Visualised the predictions on an interactive map!

# ## Citing this Notebook

# Please see [CITATION.cff](https://github.com/eds-book-gallery/94486a7f-e046-461f-bbb9-334ec7b57040/blob/main/CITATION.cff) for the full citation information. The citation file can be exported to APA or BibTex formats (learn more [here](https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-citation-files)).

# ## Additional information
# **Codebase**: version 1.0.0 with commit <mark>[16a5a1c](https://github.com/shmh40/detectreeRGB/compare/main...dev)</mark>
# 
# **License**: The code in this notebook is licensed under the MIT License. The Environmental Data Science book is licensed under the Creative Commons by Attribution 4.0 license. See further details [here](https://github.com/alan-turing-institute/environmental-ds-book/blob/master/LICENSE.md).
# 
# **Contact**: If you have any suggestion or report an issue with this notebook, feel free to [create an issue](https://github.com/alan-turing-institute/environmental-ds-book/issues/new/choose) or send a direct message to [environmental.ds.book@gmail.com](mailto:environmental.ds.book@gmail.com).

# In[ ]:


from datetime import date

print('Notebook repository version: v1.0.4')
print(f'Last tested: {date.today()}')


# ## Outputs registration
# The cell below is dedicated to save the notebook outputs for registering them into a Zenodo repository curated by the Environmental DS book.

# In[ ]:


outputs = {
    'static_figures': {
        'filenames': ['static_input','static_predictions'],
        'data':[plot_input, plot_predictions]},
    'interactive_figures': {
        'filenames': ['interactive_vector','interactive_predictions'],
        'data':[plot_vector, plot_predictions_interactive]},
}

#save static figures
if len(outputs['static_figures']['filenames']) > 0:
    [data.savefig(os.path.join(notebook_folder,outputs['static_figures']['filenames'][x]  + '.png')) for x, data in enumerate(outputs['static_figures']['data'])]

#save interactive figures
if len(outputs['interactive_figures']['filenames']) > 0:
    [hvplot.save(data, os.path.join(notebook_folder,outputs['interactive_figures']['filenames'][x]  + '.html')) for x, data in enumerate(outputs['interactive_figures']['data'])]