mask.py 2.21 KB
Newer Older
Benjamin Pillot's avatar
Benjamin Pillot committed
1
2
3
4
5
6
# -*- coding: utf-8 -*-

""" Module summary description.

More detailed description.
"""
Benjamin Pillot's avatar
Benjamin Pillot committed
7
from functools import partial
Benjamin Pillot's avatar
Benjamin Pillot committed
8
9
10
11
12
13
14

from pyrasta.io_ import ESRI_DRIVER
from pyrasta.io_.files import ShapeTempFile, _copy_to_file, NamedTempFile

import gdal


Benjamin Pillot's avatar
Benjamin Pillot committed
15
16
17
18
def _mask(arrays, no_data):
    src = arrays[0]
    mask = arrays[1]
    src[mask == 1] = no_data
Benjamin Pillot's avatar
Benjamin Pillot committed
19

Benjamin Pillot's avatar
Benjamin Pillot committed
20
    return src
Benjamin Pillot's avatar
Benjamin Pillot committed
21
22


Benjamin Pillot's avatar
Benjamin Pillot committed
23
def _raster_mask(raster, geodataframe, driver, output_type, no_data, all_touched,
24
                 window_size):
Benjamin Pillot's avatar
Benjamin Pillot committed
25
    """ Apply mask into raster
Benjamin Pillot's avatar
Benjamin Pillot committed
26

Benjamin Pillot's avatar
Benjamin Pillot committed
27
28
29
30
31
32
33
34
35
36
    """
    mask = raster.__class__.rasterize(geodataframe,
                                      raster.crs.to_wkt(),
                                      raster.x_size,
                                      raster.y_size,
                                      raster.geo_transform,
                                      burn_values=[1],
                                      all_touched=all_touched)

    return raster.__class__.raster_calculation([raster, mask],
Benjamin Pillot's avatar
Benjamin Pillot committed
37
                                               partial(_mask, no_data=no_data),
Benjamin Pillot's avatar
Benjamin Pillot committed
38
                                               gdal_driver=driver,
39
                                               output_type=raster.data_type,
Benjamin Pillot's avatar
Benjamin Pillot committed
40
                                               no_data=no_data,
Benjamin Pillot's avatar
Benjamin Pillot committed
41
42
                                               description="Compute mask",
                                               window_size=window_size,
43
44
                                               nb_processes=1,
                                               chunksize=1)
Benjamin Pillot's avatar
Benjamin Pillot committed
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65


    # with ShapeTempFile() as shp_file, \
    #         NamedTempFile(raster._gdal_driver.GetMetadata()['DMD_EXTENSION']) as out_file:
    #
    #     _copy_to_file(raster, out_file.path)
    #     geodataframe.to_file(shp_file.path, driver=ESRI_DRIVER)
    #     out_ds = gdal.Open(out_file.path, 1)
    #     gdal.Rasterize(out_ds,
    #                    shp_file.path,
    #                    bands=[bd + 1 for bd in range(raster.nb_band)],
    #                    burnValues=[10],
    #                    allTouched=all_touched)
    #
    # out_ds = None

    # return raster.__class__(out_file.path)
    # new_raster = raster.__class__(out_file.path)
    # new_raster._temp_file = out_file

    # return new_raster