Python features.rasterize函数代码示例

def test_rasterize_invalid_shapes():
    """Invalid shapes should raise an exception rather than be skipped."""
    with rasterio.Env():
        with pytest.raises(ValueError) as ex:
            rasterize([{'foo': 'bar'}], out_shape=DEFAULT_SHAPE)

        assert 'Invalid geometry object' in str(ex.value)

def test_rasterize_fill_value_dtype_mismatch(basic_geometry):
    """A fill value that doesn't match dtype should fail."""
    with pytest.raises(ValueError):
        rasterize(
            [basic_geometry], out_shape=DEFAULT_SHAPE, fill=1000000,
            default_value=2, dtype=np.uint8
        )

def test_rasterize_invalid_out_dtype(basic_geometry):
    """ A non-supported data type for out should raise an exception """

    out = np.zeros(DEFAULT_SHAPE, dtype=np.int64)
    with Env():
        with pytest.raises(ValueError):
            rasterize([basic_geometry], out=out)

def test_rasterize_missing_shapes():
    """Shapes are required for this operation."""
    with rasterio.Env():
        with pytest.raises(ValueError) as ex:
            rasterize([], out_shape=DEFAULT_SHAPE)

        assert 'No valid geometry objects' in str(ex.value)

def test_rasterize_shapes_out_dtype_mismatch(basic_geometry):
    """ Shape values must be able to fit in data type for out """

    out = np.zeros(DEFAULT_SHAPE, dtype=np.uint8)
    with Env():
        with pytest.raises(ValueError):
            rasterize([(basic_geometry, 10000000)], out=out)

def test_rasterize_supported_dtype(basic_geometry):
    """ Supported data types should return valid results """

    with Env():
        supported_types = (
            ('int16', -32768),
            ('int32', -2147483648),
            ('uint8', 255),
            ('uint16', 65535),
            ('uint32', 4294967295),
            ('float32', 1.434532),
            ('float64', -98332.133422114)
        )

        for dtype, default_value in supported_types:
            truth = np.zeros(DEFAULT_SHAPE, dtype=dtype)
            truth[2:4, 2:4] = default_value

            result = rasterize(
                [basic_geometry],
                out_shape=DEFAULT_SHAPE,
                default_value=default_value,
                dtype=dtype
            )
            assert np.array_equal(result, truth)
            assert np.dtype(result.dtype) == np.dtype(truth.dtype)

            result = rasterize(
                [(basic_geometry, default_value)],
                out_shape=DEFAULT_SHAPE
            )
            if np.dtype(dtype).kind == 'f':
                assert np.allclose(result, truth)
            else:
                assert np.array_equal(result, truth)

def test_rasterize_invalid_fill_value(basic_geometry):
    """A fill value that requires an int64 should raise an exception."""
    with pytest.raises(ValueError):
        rasterize(
            [basic_geometry], out_shape=DEFAULT_SHAPE, fill=1000000000000,
            default_value=2
        )

def test_rasterize_unsupported_dtype(basic_geometry):
    """ Unsupported types should all raise exceptions """

    with Env():
        unsupported_types = (
            ('int8', -127),
            ('int64', 20439845334323),
            ('float16', -9343.232)
        )

        for dtype, default_value in unsupported_types:
            with pytest.raises(ValueError):
                rasterize(
                    [basic_geometry],
                    out_shape=DEFAULT_SHAPE,
                    default_value=default_value,
                    dtype=dtype
                )

            with pytest.raises(ValueError):
                rasterize(
                    [(basic_geometry, default_value)],
                    out_shape=DEFAULT_SHAPE,
                    dtype=dtype
                )

def test_rasterize_invalid_value(basic_geometry):
    """A shape value that requires an int64 should raise an exception."""
    with pytest.raises(ValueError) as ex:
        rasterize(
            [(basic_geometry, 1000000000000)], out_shape=DEFAULT_SHAPE
        )

    assert 'dtype must be one of' in str(ex.value)

def test_rasterize_invalid_out_shape(basic_geometry):
    """output array shape must be 2D."""
    with pytest.raises(ValueError) as ex:
        rasterize([basic_geometry], out_shape=(1, 10, 10))
    assert 'Invalid out_shape' in str(ex.value)

    with pytest.raises(ValueError) as ex:
        rasterize([basic_geometry], out_shape=(10,))
    assert 'Invalid out_shape' in str(ex.value)

def test_rasterize_invalid_default_value(basic_geometry):
    """ A default value that requires an int64 should raise an exception """

    with Env():
        with pytest.raises(ValueError):
            rasterize(
                [basic_geometry], out_shape=DEFAULT_SHAPE,
                default_value=1000000000000
            )

def test_rasterize(basic_geometry, basic_image_2x2):
    """Rasterize operation should succeed for both an out_shape and out."""
    assert np.array_equal(
        basic_image_2x2,
        rasterize([basic_geometry], out_shape=DEFAULT_SHAPE)
    )

    out = np.zeros(DEFAULT_SHAPE)
    rasterize([basic_geometry], out=out)
    assert np.array_equal(basic_image_2x2, out)

    def set_roi(self, shape=None, geometry=None, crs=wgs84, grid=None,
                corners=None, noerase=False):
        """Set a region of interest for the dataset.
        If set succesfully, a ROI is simply a mask of the same size as the
        dataset's grid, obtained with the .roi attribute.
        I haven't decided yet if the data should be masekd out when a ROI
        has been set.
        Parameters
        ----------
        shape: path to a shapefile
        geometry: a shapely geometry
        crs: the crs of the geometry
        grid: a Grid object
        corners: a ((x0, y0), (x1, y1)) tuple of the corners of the square
        to subset the dataset to. The coordinates are not expressed in
        wgs84, set the crs keyword
        noerase: set to true in order to add the new ROI to the previous one
        """

        # The rois are always defined on the original grids, but of course
        # we take that into account when a subset is set (see roi
        # decorator below)
        ogrid = self._ogrid

        # Initial mask
        if noerase and (self.roi is not None):
            mask = self.roi
        else:
            mask = np.zeros((ogrid.ny, ogrid.nx), dtype=np.int16)

        # Several cases
        if shape is not None:
            gdf = sio.read_shapefile(shape)
            gis.transform_geopandas(gdf, to_crs=ogrid.corner_grid,
                                    inplace=True)
            with rasterio.Env():
                mask = features.rasterize(gdf.geometry, out=mask)
        if geometry is not None:
            geom = gis.transform_geometry(geometry, crs=crs,
                                          to_crs=ogrid.corner_grid)
            with rasterio.Env():
                mask = features.rasterize(np.atleast_1d(geom), out=mask)
        if grid is not None:
            _tmp = np.ones((grid.ny, grid.nx), dtype=np.int16)
            mask = ogrid.map_gridded_data(_tmp, grid, out=mask).filled(0)
        if corners is not None:
            cgrid = self._ogrid.center_grid
            xy0, xy1 = corners
            x0, y0 = cgrid.transform(*xy0, crs=crs, nearest=True)
            x1, y1 = cgrid.transform(*xy1, crs=crs, nearest=True)
            mask[np.min([y0, y1]):np.max([y0, y1])+1,
                 np.min([x0, x1]):np.max([x0, x1])+1] = 1

        self.roi = mask

def test_rasterize(basic_geometry, basic_image_2x2):
    """ Rasterize operation should succeed for both an out_shape and out """

    with rasterio.drivers():
        assert numpy.array_equal(
            basic_image_2x2,
            rasterize([basic_geometry], out_shape=DEFAULT_SHAPE)
        )

        out = numpy.zeros(DEFAULT_SHAPE)
        rasterize([basic_geometry], out=out)
        assert numpy.array_equal(basic_image_2x2, out)

def gdf_to_rst(gdf, trs, w, h, path_to_desc):
    '''
    Convert a view of a gdf to a color-coded numpy array.
    '''
    unitcolor = generate_unitcolor_lookup(path_to_desc)
    rz = rasterize([(x.geometry, unitcolor.R[gdf.mapunit[i]]) for i, x in gdf.iterrows()],
                   out_shape=(w, h), transform=trs)
    gz = rasterize([(x.geometry, unitcolor.G[gdf.mapunit[i]]) for i, x in gdf.iterrows()],
                   out_shape=(w, h), transform=trs)
    bz = rasterize([(x.geometry, unitcolor.B[gdf.mapunit[i]]) for i, x in gdf.iterrows()],
                   out_shape=(w, h), transform=trs)
    
    return np.dstack((rz, gz, bz))