Source code for rios.rat


"""
This module contains routines for reading and writing Raster
Attribute Tables (RATs). These are designed to be able to 
be called from outside of RIOS.

Within RIOS, these are called from the :class:`rios.readerinfo.ReaderInfo` 
and :class:`rios.imagewriter.ImageWriter` classes.

**Note:** It is recommended that the newer :mod:`rios.ratapplier` module be used instead of this
interface for large tables where possible.
"""
# This file is part of RIOS - Raster I/O Simplification
# Copyright (C) 2012  Sam Gillingham, Neil Flood
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import sys
import os
import warnings
from osgeo import gdal
import numpy
from . import rioserrors

if sys.version_info[0] > 2:
    # hack for Python 3 which uses str instead of basestring
    # we just use basestring
    basestring = str
    
# For supporting the automatic color table generation thing which Sam loves. 
DEFAULT_AUTOCOLORTABLETYPE = os.getenv('RIOS_DFLT_AUTOCOLORTABLETYPE', default=None)

# longer tables than this we point user to rios.colortable
MAX_RECOMMENDED_COLOR_TABLE = 256


[docs]def isColorColFromUsage(usage): "Tells if usage is one of the color column types" colorCol = (usage == gdal.GFU_Red or usage == gdal.GFU_Green or usage == gdal.GFU_Blue or usage == gdal.GFU_Alpha) return colorCol
[docs]def readColumnFromBand(gdalBand, colName): """ Given a GDAL Band, extract the Raster Attribute with the given name. Returns an array of ints or floats for numeric data types, or a list of strings. """ # get the RAT for this band rat = gdalBand.GetDefaultRAT() # get the size of the RAT numCols = rat.GetColumnCount() # if this is still None at the end # we didn't find the column colArray = None # loop thru the columns looking for the right one for col in range(numCols): if rat.GetNameOfCol(col) == colName: # found it - read the values colArray = rat.ReadAsArray(col) # one last little hack - if is a colour column, but type # was float, multiply by 255. This is so that HFA etc that stores values # between 0 and 1 is consistant with its color table and other formats usage = rat.GetUsageOfCol(col) if isColorColFromUsage(usage) and rat.GetTypeOfCol(col) == gdal.GFT_Real: colArray = colArray * 255 colArray = colArray.astype(int) # exit loop break # couldn't find named column - raise exception if colArray is None: msg = "Unable to find column named '%s'" % colName raise rioserrors.AttributeTableColumnError(msg) # return the lookup array to the caller return colArray
[docs]def readColumn(imgFile, colName, bandNumber=1): """ Given either an open gdal dataset, or a filename, extract the Raster Attribute with the given name. Returns an array of ints or floats for numeric data types, or a list of strings. """ if isinstance(imgFile, basestring): ds = gdal.Open(str(imgFile)) elif isinstance(imgFile, gdal.Dataset): ds = imgFile gdalBand = ds.GetRasterBand(bandNumber) return readColumnFromBand(gdalBand, colName)
[docs]def getColumnNamesFromBand(gdalBand): """ Return the names of the columns in the attribute table associated with the gdalBand as a list. """ # get the RAT for this band rat = gdalBand.GetDefaultRAT() colNames = [] numCols = rat.GetColumnCount() for col in range(numCols): name = rat.GetNameOfCol(col) colNames.append(name) return colNames
[docs]def getColumnNames(imgFile, bandNumber=1): """ Given either an open gdal dataset, or a filename, Return the names of the columns in the attribute table associated with the file as a list. """ if isinstance(imgFile, basestring): ds = gdal.Open(str(imgFile)) elif isinstance(imgFile, gdal.Dataset): ds = imgFile gdalBand = ds.GetRasterBand(bandNumber) return getColumnNamesFromBand(gdalBand)
[docs]def inferColumnType(sequence): """ Infer from the type of the first element in the sequence """ if isinstance(sequence[0], int) or isinstance(sequence[0], numpy.integer): colType = gdal.GFT_Integer elif isinstance(sequence[0], float) or isinstance(sequence[0], numpy.floating): colType = gdal.GFT_Real elif isinstance(sequence[0], basestring) or isinstance(sequence[0], bytes): colType = gdal.GFT_String else: colType = None return colType
[docs]def writeColumnToBand(gdalBand, colName, sequence, colType=None, colUsage=gdal.GFU_Generic): """ Given a GDAL band, Writes the data specified in sequence (can be list, tuple or array etc) to the named column in the attribute table assocated with the gdalBand. colType must be one of gdal.GFT_Integer,gdal.GFT_Real,gdal.GFT_String. can specify one of the gdal.GFU_* constants for colUsage - default is 'generic' GDAL dataset must have been created, or opened with GA_Update """ if colType is None: colType = inferColumnType(sequence) if colType is None: msg = "Can't infer type of column for sequence of %s" % type(sequence[0]) raise rioserrors.AttributeTableTypeError(msg) # check it is acually a valid type elif colType not in (gdal.GFT_Integer, gdal.GFT_Real, gdal.GFT_String): msg = "coltype must be a valid gdal column type" raise rioserrors.AttributeTableTypeError(msg) attrTbl = gdalBand.GetDefaultRAT() if attrTbl is None: # some formats eg ENVI return None # here so we need to be able to cope attrTbl = gdal.RasterAttributeTable() isFileRAT = False else: isFileRAT = True # but if it doesn't support dynamic writing # we still ahve to call SetDefaultRAT if not attrTbl.ChangesAreWrittenToFile(): isFileRAT = False # We need to ensure colname doesn't already exist colExists = False for n in range(attrTbl.GetColumnCount()): if attrTbl.GetNameOfCol(n) == colName: colExists = True colNum = n break if not colExists: # preserve usage attrTbl.CreateColumn(colName, colType, colUsage) colNum = attrTbl.GetColumnCount() - 1 rowsToAdd = len(sequence) # Imagine has trouble if not 256 items for byte if gdalBand.DataType == gdal.GDT_Byte: rowsToAdd = 256 # another hack to hide float (0-1) and int (0-255) # color table handling. # we assume that the column has already been created # of the right type appropriate for the format (maybe by calcstats) usage = attrTbl.GetUsageOfCol(colNum) if (isColorColFromUsage(usage) and attrTbl.GetTypeOfCol(colNum) == gdal.GFT_Real and colType == gdal.GFT_Integer): sequence = numpy.array(sequence, dtype=numpy.float) sequence = sequence / 255.0 attrTbl.SetRowCount(rowsToAdd) attrTbl.WriteArray(sequence, colNum) if not isFileRAT: # assume existing bands re-written # Use GDAL's exceptions to trap the error message which arises when # writing to a format which does not support it usingExceptions = gdal.GetUseExceptions() gdal.UseExceptions() try: gdalBand.SetDefaultRAT(attrTbl) except Exception: pass if not usingExceptions: gdal.DontUseExceptions()
[docs]def writeColumn(imgFile, colName, sequence, colType=None, bandNumber=1, colUsage=gdal.GFU_Generic): """ Given either an open gdal dataset, or a filename, writes the data specified in sequence (can be list, tuple or array etc) to the named column in the attribute table assocated with the file. colType must be one of gdal.GFT_Integer,gdal.GFT_Real,gdal.GFT_String. can specify one of the gdal.GFU_* constants for colUsage - default is 'generic' """ if isinstance(imgFile, basestring): ds = gdal.Open(str(imgFile), gdal.GA_Update) elif isinstance(imgFile, gdal.Dataset): ds = imgFile gdalBand = ds.GetRasterBand(bandNumber) writeColumnToBand(gdalBand, colName, sequence, colType, colUsage)
[docs]def getUsageOfColumnFromBand(gdalBand, colName): """ Given a gdalBand returns the usage of the named column """ # get the RAT for this band rat = gdalBand.GetDefaultRAT() # get the size of the RAT numCols = rat.GetColumnCount() # loop thru the columns looking for the right one for col in range(numCols): if rat.GetNameOfCol(col) == colName: return rat.GetUsageOfCol(col) msg = "Unable to find column named '%s'" % colName raise rioserrors.AttributeTableColumnError(msg)
[docs]def getUsageOfColumn(imgFile, colName, bandNumber=1): """ Given either an open gdal dataset, or a filename, returns the 'usage' of the column which can then be passed to writeColumn to preserve usage when copying """ if isinstance(imgFile, basestring): ds = gdal.Open(str(imgFile), gdal.GA_Update) elif isinstance(imgFile, gdal.Dataset): ds = imgFile gdalBand = ds.GetRasterBand(bandNumber) return getUsageOfColumnFromBand(gdalBand, colName)
[docs]def getColorTable(imgFile, bandNumber=1): """ Given either an open gdal dataset, or a filename, reads the color table as an array that can be passed to ImageWriter.setColorTable() or rat.setColorTable() The returned colour table is a numpy array, described in detail in the docstring for rat.setColorTable(). """ if isinstance(imgFile, basestring): ds = gdal.Open(str(imgFile)) elif isinstance(imgFile, gdal.Dataset): ds = imgFile gdalBand = ds.GetRasterBand(bandNumber) colorTable = gdalBand.GetColorTable() if colorTable is None: raise rioserrors.AttributeTableColumnError("Image has no color table") count = colorTable.GetCount() # count could be any size so we have to go with int colorArray = numpy.zeros((count, 5), dtype=numpy.int) for index in range(count): colorEntry = colorTable.GetColorEntry(index) arrayEntry = [index] + list(colorEntry) colorArray[index] = numpy.array(arrayEntry) return colorArray
[docs]def setColorTable(imgfile, colorTblArray, layernum=1): """ Set the color table for the specified band. You can specify either the imgfile as either a filename string or a gdal.Dataset object. The layer number defaults to 1, i.e. the first layer in the file. The color table is given as a numpy array of 5 columns. There is one row (i.e. first array index) for every value to be set, and the columns are: * pixelValue * Red * Green * Blue * Opacity The Red/Green/Blue values are on the range 0-255, with 255 meaning full color, and the opacity is in the range 0-255, with 255 meaning fully opaque. The pixels values in the first column must be in ascending order, but do not need to be a complete set (i.e. you don't need to supply a color for every possible pixel value - any not given will default to transparent black). It does not even need to be contiguous. For reasons of backwards compatability, a 4-column array will also be accepted, and will be treated as though the row index corresponds to the pixelValue (i.e. starting at zero). """ arrayShape = colorTblArray.shape if len(arrayShape) != 2: raise rioserrors.ArrayShapeError("ColorTableArray must be 2D. Found shape %s instead"%arrayShape) (numRows, numCols) = arrayShape if numRows > MAX_RECOMMENDED_COLOR_TABLE: msg = "We recommend using the rios.colortable for large color tables. " msg += "Large tables may be very slow with this function" warnings.warn(msg, DeprecationWarning) # Handle the backwards-compatible case of a 4-column array if numCols == 4: pixVals = numpy.arange(numRows) colorTbl4cols = colorTblArray elif numCols == 5: pixVals = colorTblArray[:, 0] colorTbl4cols = colorTblArray[:, 1:] else: raise rioserrors.ArrayShapeError("Color table array has %d columns, expecting 4 or 5"%numCols) # Open the image file and get the band object if isinstance(imgfile, gdal.Dataset): ds = imgfile elif isinstance(imgfile, basestring): ds = gdal.Open(imgfile, gdal.GA_Update) bandobj = ds.GetRasterBand(layernum) clrTbl = gdal.ColorTable() maxPixVal = pixVals.max() i = 0 # This loop sets an entry for every pixel value up to the largest given. Imagine # bitches if we don't do this. tblMaxVal = maxPixVal if bandobj.DataType == gdal.GDT_Byte: # For Byte files, we always add rows for entries up to 255. Imagine gets # confused if we don't tblMaxVal = 255 for pixVal in range(tblMaxVal + 1): while i < numRows and pixVals[i] < pixVal: i += 1 if i < numRows: tblPixVal = pixVals[i] if tblPixVal == pixVal: colEntry = tuple(colorTbl4cols[i]) else: colEntry = (0, 0, 0, 0) else: colEntry = (0, 0, 0, 0) clrTbl.SetColorEntry(pixVal, colEntry) bandobj.SetRasterColorTable(clrTbl)
[docs]def genColorTable(numEntries, colortype): """ Generate a colour table array. The type of colour table generated is controlled by the colortype string. Possible values are: * "rainbow" * "gray" * "random" See corresponding genColorTable_<colortype> function for details of each. """ if colortype == "rainbow": clrTbl = genColorTable_rainbow(numEntries) elif colortype == "gray": clrTbl = genColorTable_gray(numEntries) elif colortype == "random": clrTbl = genColorTable_random(numEntries) else: raise rioserrors.ColorTableGenerationError("Unknown colortype '{}'".format(colortype)) return clrTbl
[docs]def genColorTable_random(numEntries): """ Generate a color table array with the given number of entries by assigning random red/green/blue values. No attempt is made to always generate unique colours, i.e. it is randomly possibly for different entries to have the same colour. Returns the array as a 4 column array, suitable for use with the :func:`rios.rat.setColorTable` function. """ colorTable = numpy.zeros((numEntries, 4), dtype=numpy.uint8) # RGB entries, random numbers. colorTable[:, :3] = numpy.random.random_integers(0, 255, size=(numEntries, 3)) # Opacity colorTable[:, 3] = 255 return colorTable
[docs]def genColorTable_rainbow(numEntries): """ Generate a color table array with the given number of entries, with colors notionally describing a rainbow (i.e. red-orange-yellow-green-blue-indigo-violet). Probably not what a painter would call a rainbow, but it will do. Returns the array as a 4 column array, suitable for use with the :func:`rios.rat.setColorTable` function. """ colorTable = numpy.zeros((numEntries, 4), dtype=numpy.uint32) # RGB entries. Start at red, blend through to green, then through to blue, in linear steps midNdx = numEntries // 2 # Red to green colorTable[0:midNdx, 0] = numpy.mgrid[255:0:midNdx * 1j].astype(numpy.uint8) colorTable[0:midNdx, 1] = numpy.mgrid[0:255:midNdx * 1j].astype(numpy.uint8) # Green to blue colorTable[midNdx:, 1] = numpy.mgrid[255:0:(numEntries - midNdx) * 1j].astype(numpy.uint8) colorTable[midNdx:, 2] = numpy.mgrid[0:255:(numEntries - midNdx) * 1j].astype(numpy.uint8) # Opacity colorTable[:, 3] = 255 return colorTable
[docs]def genColorTable_gray(numEntries): """ Generate a color table array with the given number of entries, with all colors being shades of grey. First entry is black, last entry is white. Returns the array as a 4 column array, suitable for use with the :func:`rios.rat.setColorTable` function. """ colorTable = numpy.zeros((numEntries, 4), dtype=numpy.uint8) # RGB entries. Start at black, blend through to white, in linear steps colorTable[:, 0] = numpy.mgrid[0:255:numEntries * 1j].astype(numpy.uint8) colorTable[:, 1] = colorTable[:, 0] colorTable[:, 2] = colorTable[:, 0] # Opacity colorTable[:, 3] = 255 return colorTable