QGIS/python/plugins/processing/tools/raster.py

# -*- coding: utf-8 -*-

"""
***************************************************************************
    raster.py
    ---------------------
    Date                 : February 2013
    Copyright            : (C) 2013 by Victor Olaya  and Alexander Bruy
    Email                : volayaf at gmail dot com
***************************************************************************
*                                                                         *
*   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 2 of the License, or     *
*   (at your option) any later version.                                   *
*                                                                         *
***************************************************************************
"""

__author__ = 'Victor Olaya  and Alexander Bruy'
__date__ = 'February 2013'
__copyright__ = '(C) 2013, Victor Olaya  and Alexander Bruy'

# This will get replaced with a git SHA1 when you do a git archive

__revision__ = '$Format:%H$'

import struct
import numpy
from osgeo import gdal
from osgeo.gdalconst import *
from osgeo import osr
from PyQt4.QtCore import *
from qgis.core import *



def scanraster(layer, progress):
    filename = unicode(layer.source())
    dataset = gdal.Open(filename, GA_ReadOnly)
    band = dataset.GetRasterBand(1)
    nodata = band.GetNoDataValue()
    for y in xrange(band.YSize):
        progress.setPercentage(y / float(band.YSize) * 100)
        scanline = band.ReadRaster(0, y, band.XSize, 1, band.XSize, 1,
                                   band.DataType)
        values = struct.unpack('f' * band.XSize, scanline)
        for value in values:
            if value == nodata:
                value = None
            yield value


def mapToPixel(mX, mY, geoTransform):
    """Convert map coordinates to pixel coordinates.

    @param mX              Input map X coordinate (double)
    @param mY              Input map Y coordinate (double)
    @param geoTransform    Input geotransform (six doubles)
    @return pX, pY         Output coordinates (two doubles)
    """

    if geoTransform[2] + geoTransform[4] == 0:
        pX = (mX - geoTransform[0]) / geoTransform[1]
        pY = (mY - geoTransform[3]) / geoTransform[5]
    else:
        (pX, pY) = applyGeoTransform(mX, mY, invertGeoTransform(geoTransform))
    return (int(pX), int(pY))


def pixelToMap(pX, pY, geoTransform):
    """Convert pixel coordinates to map coordinates.

    @param pX              Input pixel X coordinate (double)
    @param pY              Input pixel Y coordinate (double)
    @param geoTransform    Input geotransform (six doubles)
    @return mX, mY         Output coordinates (two doubles)
    """

    (mX, mY) = applyGeoTransform(pX + 0.5, pY + 0.5, geoTransform)
    return (mX, mY)


def applyGeoTransform(inX, inY, geoTransform):
    """Apply a geotransform to coordinates.

    @param inX             Input coordinate (double)
    @param inY             Input coordinate (double)
    @param geoTransform    Input geotransform (six doubles)
    @return outX, outY     Output coordinates (two doubles)
    """

    outX = geoTransform[0] + inX * geoTransform[1] + inY * geoTransform[2]
    outY = geoTransform[3] + inX * geoTransform[4] + inY * geoTransform[5]
    return (outX, outY)


def invertGeoTransform(geoTransform):
    """Invert standard 3x2 set of geotransform coefficients.

    @param geoTransform        Input GeoTransform (six doubles - unaltered)
    @return outGeoTransform    Output GeoTransform (six doubles - updated)
                               on success, None if the equation is uninvertable
    """

    # We assume a 3rd row that is [1 0 0]
    # Compute determinate
    det = geoTransform[1] * geoTransform[5] - geoTransform[2] * geoTransform[4]

    if abs(det) < 0.000000000000001:
        return

    invDet = 1.0 / det

    # Compute adjoint and divide by determinate
    outGeoTransform = [0, 0, 0, 0, 0, 0]
    outGeoTransform[1] = geoTransform[5] * invDet
    outGeoTransform[4] = -geoTransform[4] * invDet

    outGeoTransform[2] = -geoTransform[2] * invDet
    outGeoTransform[5] = geoTransform[1] * invDet

    outGeoTransform[0] = (geoTransform[2] * geoTransform[3] - geoTransform[0]
                          * geoTransform[5]) * invDet
    outGeoTransform[3] = (-geoTransform[1] * geoTransform[3] + geoTransform[0]
                          * geoTransform[4]) * invDet

    return outGeoTransform


class RasterWriter:

    NODATA = -99999.0

    def __init__(self, fileName, minx, miny, maxx, maxy, cellsize,
                 nbands, crs):
        self.fileName = fileName
        self.nx = int((maxx - minx) / float(cellsize))
        self.ny = int((maxy - miny) / float(cellsize))
        self.nbands = nbands
        self.matrix = numpy.ones(shape=(self.ny, self.nx), dtype=numpy.float32)
        self.matrix[:] = self.NODATA
        self.cellsize = cellsize
        self.crs = crs
        self.minx = minx
        self.maxy = maxy

    def setValue(self, value, x, y, band=0):
        try:
            self.matrix[y, x] = value
        except IndexError:
            pass

    def getValue(self, x, y, band=0):
        try:
            return self.matrix[y, x]
        except IndexError:
            return self.NODATA

    def close(self):
        format = 'GTiff'
        driver = gdal.GetDriverByName(format)
        dst_ds = driver.Create(self.fileName, self.nx, self.ny, 1,
                               gdal.GDT_Float32)
        dst_ds.SetProjection(str(self.crs.toWkt()))
        dst_ds.SetGeoTransform([self.minx, self.cellsize, 0,
                                self.maxy, self.cellsize, 0])
        dst_ds.GetRasterBand(1).WriteArray(self.matrix)
        dst_ds = None