Python light GDAL wrapper

pyGDALite.py

from osgeo import gdal, osr, ogr
import numpy as np
from PIL import Image
from PIL import ImageOps

gdal2np_types = {gdal.GDT_UInt16: '<u2',
                 gdal.GDT_Float32: '<f4'}

np2gdal_types = {'<u2': gdal.GDT_UInt16,
		 '<f4':gdal.GDT_Float32}

class OGRVector(object):
    __class_description__ = """Class modelling a Vector layer based on OGR Class"""
    __version__ = 0.3

    def __init__(self, ds):
        self.ds = ds

    @property
    def projwkt(self):
        return self.ds.GetLayer(0).GetSpatialRef().ExportToWkt()

    @property
    def proj4(self):
        return self.ds.GetLayer(0).GetSpatialRef().ExportToProj4()

    @property
    def srs(self):
        srs = osr.SpatialReference()
        srs.ImportFromWkt(self.projwkt)

        return srs

    @property
    def points(self):
        return self.ds.GetLayer(0).GetPoints()

    @property
    def extent(self):
        return self.ds.GetLayer(0).GetExtent()

    @property
    def geometry(self):
        lyr = self.ds.GetLayer(0)
        feat = lyr.GetFeature(0)
        return feat.GetGeometryRef().Clone()

    @property
    def envelope(self):
        poly = self.geometry
        min_x, max_x, min_y, max_y = poly.GetEnvelope()
        return min_x, min_y, max_x, max_y

    @property
    def area(self):
        poly = self.geometry

        return poly.GetArea()

    def union_sq(self, other):
        poly1 = self.geometry
        poly2 = other.geometry

        min_x, max_x, min_y, max_y = poly1.Union(poly2).GetEnvelope()

        return _vector_from_polygon(_polygon_from_envelope(min_x, min_y, max_x, max_y), self.srs)

    def intersection(self, other):
        if other.proj4 != self.proj4:
            other = other.transform(self.proj4)

        poly1 = self.geometry
        poly2 = other.geometry

        return _vector_from_polygon(poly1.Intersection(poly2), self.srs)

    def intersects(self, other):
        if other.proj4 != self.proj4:
            other = other.transform(self.proj4)

        poly1 = self.geometry
        poly2 = other.geometry

        return poly1.Intersect(poly2)

    def to_geojson(self):
        poly = self.geometry

        return poly.ExportToJson()

    def transform(self, out_proj4):
        # create Coordinate Transformation
        out_srs = osr.SpatialReference()
        out_srs.ImportFromProj4(out_proj4)
        coord_transf = osr.CoordinateTransformation(self.srs, out_srs)

        geometry = self.geometry
        geometry.Transform(coord_transf)

        return _vector_from_polygon(geometry, out_srs)

    def to_raster(self, x_pix_size, y_pix_size, dtype, nodata=None):
        minX, minY, maxX, maxY = self.envelope

        x_size = int(round((maxX - minX) / abs(x_pix_size)))
        y_size = int(round((maxY - minY) / abs(y_pix_size)))

        #print(x_size, y_size)

        mem_drv = gdal.GetDriverByName('MEM')
        ds = mem_drv.Create('', x_size, y_size, 1, dtype)
        ds.SetProjection(self.projwkt)

        geot = [None, x_pix_size, 0.0, None, 0.0, y_pix_size]

        if x_pix_size >= 0:
            geot[0] = minX
        else:
            geot[0] = maxX

        if y_pix_size >= 0:
            geot[3] = minY
        else:
            geot[3] = maxY

        ds.SetGeoTransform(geot)

        # TODO Fix nodata value
        #ds.GetRasterBand(1).SetNoDataValue(nodata)

        return GDALRaster(ds)


def _vector_from_polygon(poly, srs):

    outdriver = ogr.GetDriverByName('MEMORY')
    source = outdriver.CreateDataSource('memData')
    layer = source.CreateLayer("Polygon", srs=srs, geom_type=ogr.wkbPolygon)

    feature = ogr.Feature(layer.GetLayerDefn())
    feature.SetGeometry(poly)

    # Create a field in the layer...
    # idField = ogr.FieldDefn("id", ogr.OFTInteger)
    # layer.CreateField(idField)
    # feature.SetField("id", 1)

    layer.CreateFeature(feature)

    return OGRVector(source)

def _polygon_from_envelope(minX, minY, maxX, maxY):
    ring = ogr.Geometry(ogr.wkbLinearRing)
    ring.AddPoint(minX, minY)
    ring.AddPoint(maxX, minY)
    ring.AddPoint(maxX, maxY)
    ring.AddPoint(minX, maxY)

    poly = ogr.Geometry(ogr.wkbPolygon)
    poly.AddGeometry(ring)
    poly.CloseRings()

    return poly


class GDALRaster(object):
    __class_description__ = """Class modelling a Raster layer based on GDAL"""
    __version__ = 0.3

    def __init__(self, ds):
        self.ds = ds

    @property
    def dtype(self):
        return self.ds.GetRasterBand(1).DataType

    @property
    def shape(self):
        return self.ds.RasterXSize, self.ds.RasterYSize

    @property
    def nodata(self):
        return self.ds.GetRasterBand(1).GetNoDataValue()

    @property
    def geotransform(self):
        return self.ds.GetGeoTransform()

    @property
    def pixel_size(self):
        return self.ds.GetGeoTransform()[1], self.ds.GetGeoTransform()[5]


    @property
    def projwkt(self):
        return self.ds.GetProjection()

    @property
    def proj4(self):
        coord_sys = osr.SpatialReference()
        coord_sys.ImportFromWkt(self.projwkt)

        return coord_sys.ExportToProj4()

    @property
    def srs(self):
        srs = osr.SpatialReference()
        srs.ImportFromWkt(self.projwkt)

        return srs

    @property
    def corners(self):
        ul = self.project_index(0, 0)
        lr = self.project_index(self.shape[0], self.shape[1])

        return ul, lr

    @property
    def mask(self):
        data = self.ds.ReadAsArray()
        mask = np.equal(data, self.nodata)

        mem_drv = gdal.GetDriverByName('MEM')
        ds = mem_drv.Create('', self.shape[0], self.shape[1], 1, gdal.GDT_Byte)
        ds.SetProjection(self.projwkt)
        ds.SetGeoTransform(self.geotransform)
        ds.GetRasterBand(1).SetNoDataValue(0)
        ds.GetRasterBand(1).WriteArray(mask)

        return GDALRaster(mask)

    def envelope(self):
        ul, lr = self.corners

        ring = ogr.Geometry(ogr.wkbLinearRing)
        ring.AddPoint(*ul)
        ring.AddPoint(ul[0], lr[1])
        ring.AddPoint(*lr)
        ring.AddPoint(lr[0], ul[1])

        poly = ogr.Geometry(ogr.wkbPolygon)
        poly.AddGeometry(ring)
        poly.CloseRings()

        return _vector_from_polygon(poly, self.srs)

    def project_index(self, i, j):
        gt = self.geotransform
        coords = gdal.ApplyGeoTransform(gt, i, j)

        return coords

    def get_offsets(self, x, y):
        gt = self.geotransform
        # Some version of Python?
        #ok, inv_gt = gdal.InvGeoTransform(gt)
        inv_gt = gdal.InvGeoTransform(gt)
        offsets = gdal.ApplyGeoTransform(inv_gt, x, y)
        i_off, j_off = map(int, map(round, offsets))

        return i_off, j_off

    def offset_geotransform(self, i, j):
        geot = self.geotransform
        min_x = i * geot[1] + geot[0]
        min_y = j * geot[5] + geot[3]

        return min_x, geot[1], .0, min_y, .0, geot[5]

    def corners_from_envelope(self, minX, minY, maxX, maxY):
        if self.geotransform[1] > 0:
            ulX = minX
            lrX = maxX
        else:
            ulX = maxX
            lrX = minX

        if self.geotransform[5] > 0:
            ulY = minY
            lrY = maxY
        else:
            ulY = maxY
            lrY = minY

        return ulX, ulY, lrX, lrY

    def clip(self, polygon):

        if self.envelope().intersects(polygon):
            inters_poly = self.envelope().intersection(polygon)
            ul_x, ul_y, lr_x, lr_y = self.corners_from_envelope(*inters_poly.envelope)

            ul_i, ul_j = self.get_offsets(ul_x, ul_y)
            lr_i, lr_j = self.get_offsets(lr_x, lr_y)

            geot = self.offset_geotransform(ul_i, ul_j)

            mem_drv = gdal.GetDriverByName('MEM')
            ds = mem_drv.Create('', lr_i-ul_i, lr_j-ul_j, 1, self.dtype)
            ds.SetProjection(self.projwkt)
            ds.SetGeoTransform(geot)
            #ds.GetRasterBand(1).SetNoDataValue(self.nodata)

            ds.GetRasterBand(1).WriteArray(self.ds.ReadAsArray(ul_i, ul_j, lr_i-ul_i, lr_j-ul_j))

            return GDALRaster(ds)

        else:
            return None

    def merge(self, in_raster):
        # For this first release we suppose in raster fits inside the caller raster....
        # TODO --- Make it general for any size of in_raster

        trans = gdal.Transformer(in_raster.ds, self.ds, [])
        success, xyz = trans.TransformPoint(False, 0, 0)
        x, y, z = map(int, map(round, xyz))

        data = in_raster.ds.ReadAsArray()
        chunk = self.ds.ReadAsArray(x, y, data.shape[1], data.shape[0])
        m_chunk = np.ma.masked_array(chunk, mask=np.equal(chunk, 0))
        m_chunk[m_chunk.mask] = data[m_chunk.mask]

        self.ds.GetRasterBand(1).WriteArray(m_chunk.data, x, y)

    def reproject(self, polygon, x_pix_size, y_pix_size):
        subset = self.clip(polygon)
        if subset is not None:
            dest = polygon.to_raster(x_pix_size, y_pix_size, self.dtype, self.nodata)
            gdal.ReprojectImage(subset.ds, dest.ds, self.projwkt, polygon.projwkt, gdal.GRA_NearestNeighbour)
            return dest
        else:
	    return None

    def as_rgba(self, clip_value):
        from PIL import Image
        import matplotlib.pyplot as plt
        #size = (100, 100)

        array = self.ds.ReadAsArray()
        np.clip(array, 0, clip_value, out=array)
	array = array.astype(np.float32)
        array *= 1. / clip_value

        im = Image.fromarray(np.uint8(plt.cm.gist_earth(array)*255))
        im = im.convert("RGBA")
        im = ImageOps.fit(im, self.shape, Image.ANTIALIAS)
        return np.asarray(im)

def new_raster(path):
    #print path
    ds = gdal.Open(path, gdal.GA_ReadOnly)
    return GDALRaster(ds)

def _vector_from_polygon(poly, srs):
    outdriver = ogr.GetDriverByName('MEMORY')
    source = outdriver.CreateDataSource('memData')
    layer = source.CreateLayer("Polygon", srs=srs, geom_type=ogr.wkbPolygon)

    feature = ogr.Feature(layer.GetLayerDefn())
    feature.SetGeometry(poly)

    layer.CreateFeature(feature)

    return OGRVector(source)

def new_vector_from_geojson(geojson, proj4):
    srs = osr.SpatialReference()
    srs.ImportFromProj4(proj4)
    return _vector_from_polygon(ogr.CreateGeometryFromJson(geojson), srs)