Source Code for Module madrona.heatmap.models

  1  from django.contrib.gis.db import models 
  2  from django.conf import settings 
  3  from madrona.studyregion.models import * 
  4  import os 
  5  import sys 
  6  import numpy 
  7  import tempfile 
  8  from osgeo import gdal, osr 
  9   
 10  CELL_SIZE = 50 
 11   
 12 -def readArray(input): 
 13      """Borrowed from Matt Perry""" 
 14      data = gdal.Open(input) 
 15      band = data.GetRasterBand(1) 
 16   
 17      return band.ReadAsArray() 
 18   
 19 -def create_blank_raster(extent,cellsize,outfile,format,srid=settings.GEOMETRY_DB_SRID): 
 20      """ 
 21      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 
 22      to set the projection. 
 23      """ 
 24      srs = osr.SpatialReference() 
 25      srs.ImportFromEPSG(srid) 
 26      ydist = extent[3] - extent[1] 
 27      xdist = extent[2] - extent[0] 
 28      xcount = int((xdist / cellsize) + 1) 
 29      ycount = int((ydist / cellsize) + 1) 
 30   
 31      # Create the blank numpy array 
 32      outArray = numpy.zeros((ycount, xcount)) 
 33   
 34      # Create output raster   
 35      driver = gdal.GetDriverByName(format) 
 36      dst_ds = driver.Create(outfile, xcount, ycount, 1, gdal.GDT_UInt16) 
 37   
 38      # This is bizzarly complicated 
 39      # the GT(2) and GT(4) coefficients are zero,      
 40      # and the GT(1) is pixel width, and GT(5) is pixel height.      
 41      # The (GT(0),GT(3)) position is the top left corner of the top left pixel 
 42      gt = (extent[0],cellsize,0,extent[3],0,(cellsize * -1.)) 
 43      dst_ds.SetGeoTransform(gt) 
 44      dst_ds.SetProjection(srs.ExportToWkt()) 
 45   
 46      dst_band = dst_ds.GetRasterBand(1) 
 47      dst_band.WriteArray(outArray,0,0) 
 48      dst_ds = None 
 49      return outfile 
 50   
 51 -def create_raster_from_matrix(matrix,outfile,extent=None,cellsize=CELL_SIZE,format='GTiff',srid=settings.GEOMETRY_DB_SRID): 
 52      """ 
 53      Creates a raster dataset with all 1s where there are shapes in the shapefile. 
 54      """ 
 55      if extent is None: 
 56          extent = StudyRegion.objects.current().geometry.extent 
 57      srs = osr.SpatialReference() 
 58      srs.ImportFromEPSG(srid) 
 59      ydist = extent[3] - extent[1] 
 60      xdist = extent[2] - extent[0] 
 61      xcount = int((xdist / cellsize) + 1) 
 62      ycount = int((ydist / cellsize) + 1) 
 63   
 64      # Create output raster   
 65      driver = gdal.GetDriverByName(format) 
 66      dst_ds = driver.Create(outfile, xcount, ycount, 1, gdal.GDT_UInt16) 
 67   
 68      # This is bizzarly complicated 
 69      # the GT(2) and GT(4) coefficients are zero,      
 70      # and the GT(1) is pixel width, and GT(5) is pixel height.      
 71      # The (GT(0),GT(3)) position is the top left corner of the top left pixel 
 72      gt = (extent[0],cellsize,0,extent[3],0,(cellsize * -1.)) 
 73      dst_ds.SetGeoTransform(gt) 
 74      dst_ds.SetProjection(srs.ExportToWkt()) 
 75   
 76      dst_band = dst_ds.GetRasterBand(1) 
 77      dst_band.WriteArray(matrix,0,0) 
 78      dst_band.SetNoDataValue(0.0) 
 79      min_val, max_val = dst_band.ComputeRasterMinMax() 
 80      start_color = (255,255,0,160) 
 81      end_color = (255,0,0,160) 
 82      ct = gdal.ColorTable() 
 83      ct.CreateColorRamp(int(min_val),start_color,int(max_val),end_color) 
 84      dst_band.SetColorTable(ct) 
 85      dst_ds = None 
 86      return outfile 
 87   
 88 -def shapefile_to_matrix(shapefile,cellsize=CELL_SIZE,extent=None): 
 89      if extent is None: 
 90          extent = StudyRegion.objects.current().geometry.extent 
 91      tf = tempfile.mktemp() 
 92      rast = create_blank_raster(extent,cellsize,tf,'GTiff') 
 93      layer_arg = '-l ' + os.path.basename(shapefile).split(os.path.extsep)[0] 
 94      sys_string = 'gdal_rasterize -b 1 -burn 1 %s %s %s' % (layer_arg,shapefile,tf) 
 95      os.system(sys_string) 
 96      matrix = readArray(tf) 
 97      os.remove(tf) 
 98      return matrix 
 99   
100 -def create_heatmap(shp_list): 
101      matrices = [] 
102      for shp in shp_list: 
103          matrices.append(shapefile_to_matrix(shp)) 
104      for i,matrix in enumerate(matrices): 
105          if i == 0: 
106              mat = matrix 
107          else: 
108              mat += matrix 
109      outfile = tempfile.NamedTemporaryFile(suffix='.tif', mode='w+b') 
110      outfile.close() 
111      return create_raster_from_matrix(mat,outfile.name) 
112