Coverage for madrona/heatmap/models : 18%

from django.contrib.gis.db import models 

from django.conf import settings 

from madrona.studyregion.models import * 

import os 

import sys 

import numpy 

import tempfile 

from osgeo import gdal, osr 

CELL_SIZE = 50 

def readArray(input): 

    """Borrowed from Matt Perry""" 

    data = gdal.Open(input) 

    band = data.GetRasterBand(1) 

    return band.ReadAsArray() 

def create_blank_raster(extent,cellsize,outfile,format,srid=settings.GEOMETRY_DB_SRID): 

    """ 

    Creates a blank raster dataset with all zeros.  This code came from Matt Perry: http://perrygeo.googlecode.com/hg/gis-bin/blank_raster.py and got modified 

    to set the projection. 

    """ 

    srs = osr.SpatialReference() 

    srs.ImportFromEPSG(srid) 

    ydist = extent[3] - extent[1] 

    xdist = extent[2] - extent[0] 

    xcount = int((xdist / cellsize) + 1) 

    ycount = int((ydist / cellsize) + 1) 

    # Create the blank numpy array 

    outArray = numpy.zeros((ycount, xcount)) 

    # Create output raster   

    driver = gdal.GetDriverByName(format) 

    dst_ds = driver.Create(outfile, xcount, ycount, 1, gdal.GDT_UInt16) 

    # This is bizzarly complicated 

    # the GT(2) and GT(4) coefficients are zero,      

    # and the GT(1) is pixel width, and GT(5) is pixel height.      

    # The (GT(0),GT(3)) position is the top left corner of the top left pixel 

    gt = (extent[0],cellsize,0,extent[3],0,(cellsize * -1.)) 

    dst_ds.SetGeoTransform(gt) 

    dst_ds.SetProjection(srs.ExportToWkt()) 

    dst_band = dst_ds.GetRasterBand(1) 

    dst_band.WriteArray(outArray,0,0) 

    dst_ds = None 

    return outfile 

def create_raster_from_matrix(matrix,outfile,extent=None,cellsize=CELL_SIZE,format='GTiff',srid=settings.GEOMETRY_DB_SRID): 

    """ 

    Creates a raster dataset with all 1s where there are shapes in the shapefile. 

    """ 

    if extent is None: 

        extent = StudyRegion.objects.current().geometry.extent 

    srs = osr.SpatialReference() 

    srs.ImportFromEPSG(srid) 

    ydist = extent[3] - extent[1] 

    xdist = extent[2] - extent[0] 

    xcount = int((xdist / cellsize) + 1) 

    ycount = int((ydist / cellsize) + 1) 

    # Create output raster   

    driver = gdal.GetDriverByName(format) 

    dst_ds = driver.Create(outfile, xcount, ycount, 1, gdal.GDT_UInt16) 

    # This is bizzarly complicated 

    # the GT(2) and GT(4) coefficients are zero,      

    # and the GT(1) is pixel width, and GT(5) is pixel height.      

    # The (GT(0),GT(3)) position is the top left corner of the top left pixel 

    gt = (extent[0],cellsize,0,extent[3],0,(cellsize * -1.)) 

    dst_ds.SetGeoTransform(gt) 

    dst_ds.SetProjection(srs.ExportToWkt()) 

    dst_band = dst_ds.GetRasterBand(1) 

    dst_band.WriteArray(matrix,0,0) 

    dst_band.SetNoDataValue(0.0) 

    min_val, max_val = dst_band.ComputeRasterMinMax() 

    start_color = (255,255,0,160) 

    end_color = (255,0,0,160) 

    ct = gdal.ColorTable() 

    ct.CreateColorRamp(int(min_val),start_color,int(max_val),end_color) 

    dst_band.SetColorTable(ct) 

    dst_ds = None 

    return outfile 

def shapefile_to_matrix(shapefile,cellsize=CELL_SIZE,extent=None): 

    if extent is None: 

        extent = StudyRegion.objects.current().geometry.extent 

    tf = tempfile.mktemp() 

    rast = create_blank_raster(extent,cellsize,tf,'GTiff') 

    layer_arg = '-l ' + os.path.basename(shapefile).split(os.path.extsep)[0] 

    sys_string = 'gdal_rasterize -b 1 -burn 1 %s %s %s' % (layer_arg,shapefile,tf) 

    os.system(sys_string) 

    matrix = readArray(tf) 

    os.remove(tf) 

    return matrix 

def create_heatmap(shp_list): 

    matrices = [] 

    for shp in shp_list: 

        matrices.append(shapefile_to_matrix(shp)) 

    for i,matrix in enumerate(matrices): 

        if i == 0: 

            mat = matrix 

        else: 

            mat += matrix 

    outfile = tempfile.NamedTemporaryFile(suffix='.tif', mode='w+b') 

    outfile.close() 

    return create_raster_from_matrix(mat,outfile.name)