PointDatabase/read_DEM.py at master · bpjelley/PointDatabase · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Dec 22 15:35:13 2018

@author: ben
"""

from osgeo import gdal, gdalconst
import numpy as np
import matplotlib.pyplot as plt


def read_DEM(file, band_num=1, bounds=None, skip=1, asPoints=False, keepAll=False, getProjection=False):
    """
        Read a raster from a DEM file
    """
    ds=gdal.Open(file, gdalconst.GA_ReadOnly)
    if getProjection:
        proj=ds.GetProjection()
    band=ds.GetRasterBand(band_num)
    GT=ds.GetGeoTransform()
    nodataValue=band.GetNoDataValue()
    # ii and jj are the pixel center coordinates.  0,0 in GDAL is the upper-left
    # corner of the first pixel.
    ii=np.arange(0, band.XSize)+0.5
    jj=np.arange(0, band.YSize)-0.5
    x=GT[0]+GT[1]*ii
    y=GT[3]+GT[5]*jj
    if bounds is not None:
        cols = np.where(( x>=bounds[0][0] ) & ( x<= bounds[0][1] ))[0]
        rows = np.where(( y>=bounds[1][0] ) & ( y<= bounds[1][1] ))[0]
    else:
        rows=np.arange(band.YSize, dtype=int)
        cols=np.arange(band.XSize, dtype=int)

    z=band.ReadAsArray(int(cols[0]), int(rows[0]), int(cols[-1]-cols[0]+1), int(rows[-1]-rows[0]+1))
    ds=None

    if skip >1:
        z=z[::skip, ::skip]
        cols=cols[::skip]
        rows=rows[::skip]
    if nodataValue is not None and np.isfinite(nodataValue):
        bad = z==np.array(nodataValue).astype(z.dtype)
        z = np.float64(z)
        z[bad] = np.NaN
    else:
        z = np.float64(z)
    x=x[cols]
    y=y[rows]
    if asPoints:
        x,y=np.meshgrid(x, y)
        if keepAll:
            if getProjection:
                return x.ravel(), y.ravel(), z.ravel(), proj
            else:
                return x.ravel(), y.ravel(), z.ravel()
        else:
            keep=np.isfinite(z.ravel())
            if getProjection:
                return x.ravel()[keep], y.ravel()[keep], z.ravel()[keep], proj
            else:
                return x.ravel()[keep], y.ravel()[keep], z.ravel()[keep]
    else:
        if getProjection:
            return x, y[::-1], z[::-1, :], proj
        else:
            return x, y[::-1], z[::-1, :]