#!/usr/bin/env python
"""
Calculate NDVI from RAT using RIOS rat
module.

Dan Clewley

https://spectraldifferences.wordpress.com/

"""

from osgeo import gdal
from rios import rat

# Open RAT dataset
rat_dataset = gdal.Open("clumps.kea", gdal.GA_Update)

# Get columns with average red and NIR for each object
red = rat.readColumn(rat_dataset, "RedAvg")
nir = rat.readColumn(rat_dataset, "NIR1Avg")

ndvi = (nir - red) / (nir + red)

# Write out column
rat.writeColumn(rat_dataset, "NDVI", ndvi)

# Close RAT dataset
rat_dataset = None

#!/usr/bin/env python
"""
Calculate NDVI from RAT using RIOS ratapplier
module.

Dan Clewley

https://spectraldifferences.wordpress.com/

"""
from rios import ratapplier

def _ratapplier_calc_ndvi(info, inputs, outputs):
    """
    Calculate NDVI from RAT.

    Called by ratapplier
    """

    # Get columns with average red and NIR for each object
    # within block
    red = getattr(inputs.inrat, "RedAvg")
    nir = getattr(inputs.inrat, "NIR1Avg")

    # Calculate NDVI
    ndvi = (nir - red) / (nir + red)

    # Save to 'NDVI' column (will create if doesn't exist)
    setattr(outputs.outrat,"NDVI", ndvi)

if __name__ == "__main__":

    # Set up rat applier for input / output
    in_rats = ratapplier.RatAssociations()
    out_rats = ratapplier.RatAssociations()

    # Pass in clumps file
    # Same file is used for input and output to write
    # to existing RAT
    in_rats.inrat = ratapplier.RatHandle("clumps.kea")
    out_rats.outrat = ratapplier.RatHandle("clumps.kea")

    # Apply function to all rows in chunks
    ratapplier.apply(_ratapplier_calc_ndvi, in_rats, out_rats)