Skip to content

Commit

Permalink
Allow io.grid_to_raster to return a rio.MemoryFile (#257)
Browse files Browse the repository at this point in the history
  • Loading branch information
atmorling authored Aug 30, 2024
1 parent 3b66510 commit cf861ed
Show file tree
Hide file tree
Showing 4 changed files with 49 additions and 12 deletions.
2 changes: 1 addition & 1 deletion doc/source/notebooks/04. EcoMap & EcoPlot/EcoMap.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -332,7 +332,7 @@
"\n",
"m = EcoMap(width=800, height=600)\n",
"m.add_layer(EcoMap.get_named_tile_layer(\"OpenStreetMap\"))\n",
"m.add_geotiff(path=os.path.join(output_dir, \"mara_dem.tif\"), zoom=True, cmap=\"jet\")\n",
"m.add_geotiff(tiff=os.path.join(output_dir, \"mara_dem.tif\"), zoom=True, cmap=\"jet\")\n",
"m"
]
},
Expand Down
27 changes: 20 additions & 7 deletions ecoscope/io/raster.py
Original file line number Diff line number Diff line change
Expand Up @@ -208,7 +208,7 @@ def raster_to_gdf(raster_path):
)


def grid_to_raster(grid=None, val_column="", out_dir="", raster_name="grid.tif", xlen=5000, ylen=5000, grid_crs=4326):
def grid_to_raster(grid=None, val_column="", out_dir="", raster_name=None, xlen=5000, ylen=5000):
"""
Save a GeoDataFrame grid to a raster.
"""
Expand All @@ -223,16 +223,29 @@ def grid_to_raster(grid=None, val_column="", out_dir="", raster_name="grid.tif",

raster_profile = ecoscope.io.raster.RasterProfile(
pixel_size=xlen,
crs=grid_crs,
crs=grid.crs.to_epsg(),
nodata_value=np.nan,
band_count=1,
)

raster_profile.raster_extent = ecoscope.io.raster.RasterExtent(
x_min=bounds[0], x_max=bounds[2], y_min=bounds[1], y_max=bounds[3]
)
ecoscope.io.raster.RasterPy.write(
ndarray=vals,
fp=os.path.join(out_dir, raster_name),
**raster_profile,
)

if raster_name:
ecoscope.io.raster.RasterPy.write(
ndarray=vals,
fp=os.path.join(out_dir, raster_name),
**raster_profile,
)
else:
memfile = rio.MemoryFile()
memfile.open(
driver="GTiff",
width=raster_profile.pop("columns"),
height=raster_profile.pop("rows"),
count=raster_profile.pop("band_count"),
**raster_profile,
).write(vals, 1)

return memfile
9 changes: 5 additions & 4 deletions ecoscope/mapping/map.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

import numpy as np
import pandas as pd
import rasterio as rio
from io import BytesIO
from typing import Dict, IO, List, Optional, TextIO, Union
from pathlib import Path
Expand Down Expand Up @@ -368,7 +369,7 @@ def zoom_to_bounds(self, feat: Union[BaseLayer, List[BaseLayer], gpd.GeoDataFram

def add_geotiff(
self,
path: str,
tiff: str | rio.MemoryFile,
zoom: bool = False,
cmap: Union[str, mpl.colors.Colormap] = None,
opacity: float = 0.7,
Expand All @@ -380,16 +381,16 @@ def add_geotiff(
Parameters
----------
path : str
The path to the local tiff
tiff : str | rio.MemoryFile
The string path to a tiff on disk or a rio.MemoryFile
zoom : bool
Whether to zoom the map to the bounds of the tiff
cmap: str or matplotlib.colors.Colormap
The colormap to apply to the raster
opacity: float
The opacity of the overlay
"""
with rasterio.open(path) as src:
with rasterio.open(tiff) as src:
transform, width, height = rasterio.warp.calculate_default_transform(
src.crs, "EPSG:4326", src.width, src.height, *src.bounds
)
Expand Down
23 changes: 23 additions & 0 deletions tests/test_ecomap.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import os
import random
import ee
import geopandas as gpd
import pandas as pd
Expand Down Expand Up @@ -215,6 +217,27 @@ def test_add_geotiff_with_cmap():
assert isinstance(m.layers[1], BitmapLayer)


def test_add_geotiff_in_mem_with_cmap():
AOI = gpd.read_file(os.path.join("tests/sample_data/vector", "maec_4zones_UTM36S.gpkg"))

grid = gpd.GeoDataFrame(
geometry=ecoscope.base.utils.create_meshgrid(
AOI.unary_union, in_crs=AOI.crs, out_crs=AOI.crs, xlen=5000, ylen=5000, return_intersecting_only=False
)
)

grid["fake_density"] = grid.apply(lambda _: random.randint(1, 50), axis=1)
raster = ecoscope.io.raster.grid_to_raster(
grid,
val_column="fake_density",
)

m = EcoMap()
m.add_geotiff(raster, cmap="jet")
assert len(m.layers) == 2
assert isinstance(m.layers[1], BitmapLayer)


def test_add_datashader_gdf(point_gdf):
m = EcoMap()
img, bounds = datashade_gdf(point_gdf, "point")
Expand Down

0 comments on commit cf861ed

Please sign in to comment.