kopia lustrzana https://github.com/OpenDroneMap/ODM
83 wiersze
2.7 KiB
Python
83 wiersze
2.7 KiB
Python
# A script to calculate the NDVI from a color-infrared orthophoto.
|
|
# requires python-gdal
|
|
|
|
import numpy
|
|
import argparse
|
|
import os.path
|
|
try:
|
|
from osgeo import gdal
|
|
from osgeo import osr
|
|
except ImportError:
|
|
raise ImportError("You need to install python-gdal. run `apt-get install python-gdal`")
|
|
exit()
|
|
|
|
|
|
def parse_args():
|
|
p = argparse.ArgumentParser("A script that calculates the NDVI of a CIR orthophoto")
|
|
|
|
p.add_argument("orthophoto", metavar="<orthophoto.tif>",
|
|
type=argparse.FileType('r'),
|
|
help="The CIR orthophoto. Must be a GeoTiff.")
|
|
p.add_argument("nir", metavar="N", type=int,
|
|
help="NIR band number")
|
|
p.add_argument("vis", metavar="N", type=int,
|
|
help="Vis band number")
|
|
p.add_argument("out", metavar="<outfile.tif>",
|
|
type=argparse.FileType('w'),
|
|
help="The output file. Also must be in GeoTiff format")
|
|
p.add_argument("--overwrite", "-o",
|
|
action='store_true',
|
|
default=False,
|
|
help="Will overwrite output file if it exists. ")
|
|
return p.parse_args()
|
|
|
|
|
|
def calc_ndvi(nir, vis):
|
|
"""
|
|
Calculates the NDVI of an orthophoto using nir and vis bands.
|
|
:param nir: An array containing the nir band
|
|
:param vis: An array containing the vis band
|
|
:return: An array that will be exported as a tif
|
|
"""
|
|
|
|
# Take the orthophoto and do nir - vis / nir + vis
|
|
# for each cell, calculate ndvi (masking out where divide by 0)
|
|
ndvi = numpy.empty(nir.shape, dtype=float)
|
|
mask = numpy.not_equal((nirb + visb), 0.0)
|
|
return numpy.choose(mask, (-1.0, numpy.true_divide(numpy.subtract(nirb, visb), numpy.add(nirb, visb))))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
rootdir = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
# Parse args
|
|
args = parse_args()
|
|
|
|
if not args.overwrite and os.path.isfile(os.path.join(rootdir, args.out.name)):
|
|
print("File exists, rename or use -o to overwrite.")
|
|
exit()
|
|
|
|
# import raster
|
|
raster = gdal.Open(args.orthophoto.name)
|
|
orthophoto = raster.ReadAsArray()
|
|
# parse out bands
|
|
nirb = orthophoto[args.nir - 1].astype(float)
|
|
visb = orthophoto[args.vis - 1].astype(float)
|
|
|
|
outfile = args.out
|
|
|
|
# Do ndvi calc
|
|
ndvi = calc_ndvi(nirb, visb)
|
|
|
|
# export raster
|
|
out_driver = gdal.GetDriverByName('GTiff')\
|
|
.Create(outfile.name, int(ndvi.shape[1]), int(ndvi.shape[0]), 1, gdal.GDT_Float32)
|
|
outband = out_driver.GetRasterBand(1)
|
|
outband.WriteArray(ndvi)
|
|
outcrs = osr.SpatialReference()
|
|
outcrs.ImportFromWkt(raster.GetProjectionRef())
|
|
out_driver.SetProjection(outcrs.ExportToWkt())
|
|
out_driver.SetGeoTransform(raster.GetGeoTransform())
|
|
outband.FlushCache()
|