Commit 0bd605b0 authored by Jordi Inglada's avatar Jordi Inglada
Browse files

ENH: clean up, doc, refactor

parent 1af867c4
......@@ -11,10 +11,14 @@ import numpy as np
import rasterio as rio
from pyproj import CRS, Transformer
from rasterio.merge import merge
from rasterio.warp import Resampling, reproject
from shapely.geometry import Polygon
class BoundingBox:
"""
Class modeling a lat, long bounding box
"""
def __init__(self, lonmin: float, latmin: float, lonmax: float,
latmax: float):
self.lonmin = min(lonmax, lonmin)
......@@ -23,21 +27,23 @@ class BoundingBox:
self.latmax = max(latmax, latmin)
@dataclass
@dataclass(frozen=True)
class SRTMTileId:
""" Class modeling a SRTM tile """
lon: int
lat: int
def srtm_id_to_name(id: SRTMTileId) -> str:
northing = 'S' if id.lat < 0 else 'N'
easting = 'W' if id.lon < 0 else 'E'
lat = str(np.abs(id.lat)).zfill(2)
lon = str(np.abs(id.lon)).zfill(3)
return f"{northing}{lat}{easting}{lon}"
def name(self) -> str:
""" Build the name of the tile from the lon, lat values"""
northing = 'S' if self.lat < 0 else 'N'
easting = 'W' if self.lon < 0 else 'E'
lat = str(np.abs(self.lat)).zfill(2)
lon = str(np.abs(self.lon)).zfill(3)
return f"{northing}{lat}{easting}{lon}"
def srtm_tiles_from_bbox(bbox: BoundingBox) -> List[SRTMTileId]:
""" Return a list of SRTMTileId intersecting a bounding box"""
latmin = int(np.floor(bbox.latmin))
latmax = int(np.floor(bbox.latmax))
lonmin = int(np.floor(bbox.lonmin))
......@@ -48,7 +54,8 @@ def srtm_tiles_from_bbox(bbox: BoundingBox) -> List[SRTMTileId]:
]
def mgrs_polygon(tile: str) -> Polygon:
def get_polygon_mgrs_tile(tile: str) -> Polygon:
""" Get the shapely.Polygon corresponding to a MGRS tile"""
assert tile[0] != 'T'
mgrs_df = gpd.read_file(
os.path.join(os.path.dirname(os.path.abspath(__file__)),
......@@ -56,33 +63,44 @@ def mgrs_polygon(tile: str) -> Polygon:
return mgrs_df[mgrs_df.Name == tile].iloc[0].geometry
def mgrs_bbox(tile: str) -> BoundingBox:
poly = mgrs_polygon(tile)
def get_bbox_mgrs_tile(tile: str) -> BoundingBox:
""" Get a bounding box in '+proj=latlong' for a MGRS tile"""
poly = get_polygon_mgrs_tile(tile)
return BoundingBox(*poly.bounds)
def srtm_tiles_from_mgrs_tile(tile: str) -> List[SRTMTileId]:
return srtm_tiles_from_bbox(mgrs_bbox(tile))
def get_srtm_tiles_for_mgrs_tile(tile: str) -> List[SRTMTileId]:
""" Get the list of SRTM tiles intersecting a MGRS tile"""
return srtm_tiles_from_bbox(get_bbox_mgrs_tile(tile))
def mgrs_transform(tile: str) -> rio.Affine:
ul, ur, ll, lr, _ = mgrs_polygon(tile).exterior.coords
def get_transform_mgrs_tile(tile: str) -> rio.Affine:
""" Get the rasterio.Affine transform for a MGRS tile in its own CRS"""
ul, ur, ll, lr, _ = get_polygon_mgrs_tile(tile).exterior.coords
transformer = Transformer.from_crs('+proj=latlong',
crs_for_mgrs_tile(tile))
get_crs_mgrs_tile(tile))
x0, y0 = transformer.transform([ul[0]], [ul[1]])
return rio.Affine(10.0, 0.0, np.round(x0[0]), 0.0, -10.0, np.round(y0[0]))
def crs_for_mgrs_tile(tile: str) -> CRS:
def get_crs_mgrs_tile(tile: str) -> CRS:
"""Get the pyproj.CRS for a MGRS tile.
The tile is given as 31TCJ: the 2 digits correspond to the UTM zone
and the first letter is used to determine the northing (N-Z are
North hemisphere, see
https://en.wikipedia.org/wiki/Military_Grid_Reference_System#Grid_zone_designation)
"""
zone = int(tile[:2])
south = tile[2] < 'N'
return CRS.from_dict({'proj': 'utm', 'zone': zone, 'south': south})
@dataclass(frozen=True)
class DEM:
"""Class modeling a 3 channel (elevation, slope, aspect) DEM
"""
elevation: np.ndarray
slope: np.ndarray
aspect: np.ndarray
......@@ -90,8 +108,26 @@ class DEM:
transform: Optional[rio.Affine]
def as_stack(self):
"""Get the DEM as a single 3 channel numpy array"""
return np.stack([self.elevation, self.slope, self.aspect], axis=0)
def write(self, out_file: str):
"""Write the DEM to a geotiff file"""
with rio.open(out_file,
'w',
driver='GTiff',
height=self.elevation.shape[0],
width=self.elevation.shape[1],
count=3,
nodata=-32768.0,
dtype=self.elevation.dtype,
compress='lzw',
crs=self.crs,
transform=self.transform) as ds:
ds.write(self.elevation, 1)
ds.write(self.slope, 2)
ds.write(self.aspect, 3)
@dataclass(frozen=True)
class SRTM:
......@@ -116,27 +152,44 @@ class SRTM:
crs='+proj=latlong',
transform=transform)
def get_dem_for_mgrs_tile(self, tile: str) -> DEM:
"""Build a DEM (elevation, slope and aspect) for an MGRS tile with the
union of the intersecting SRTM tiles. It keeps CRS and extent of the
SRTM tiles.
def get_dem_mgrs_tile(self, tile: str) -> DEM:
"""Build a DEM (elevation, slope and aspect) for an MGRS tile
with the union of the intersecting SRTM tiles. It keeps CRS and
extent of the SRTM tiles.
"""
srtm_tiles = srtm_tiles_from_mgrs_tile(tile)
srtm_tiles = get_srtm_tiles_for_mgrs_tile(tile)
return self.get_dem_from_tiles(srtm_tiles)
def get_dem_for_bbox(self, bbox: BoundingBox) -> DEM:
"""Build a DEM (elevation, slope and aspect) for a bounding box
with the union of the intersecting SRTM tiles. It keeps CRS and extent
of the SRTM tiles.
with the union of the intersecting SRTM tiles. It keeps CRS and
extent of the SRTM tiles.
"""
srtm_tiles = srtm_tiles_from_bbox(bbox)
return self.get_dem_from_tiles(srtm_tiles)
def __build_hgt(self,
tiles: List[SRTMTileId]) -> Tuple[np.ndarray, rio.Affine]:
file_names = [
f"{self.base_dir}/{srtm_id_to_name(t)}.hgt" for t in tiles
]
file_names = [f"{self.base_dir}/{t.name()}.hgt" for t in tiles]
return merge(file_names) # type: ignore
def get_dem_mgrs_tile(tile: str) -> DEM:
""" Get a 10m resolution DEM on the geometry of a MGRS tile"""
mgrs_trsf = get_transform_mgrs_tile(tile)
mgrs_crs = get_crs_mgrs_tile(tile)
dst_dem = np.zeros((3, 10980, 10980))
dem_handler = SRTM()
srtm_dem = dem_handler.get_dem_mgrs_tile(tile)
dst_dem, dst_dem_transform = reproject(srtm_dem.as_stack(),
destination=dst_dem,
src_transform=srtm_dem.transform,
src_crs=srtm_dem.crs,
dst_transform=mgrs_trsf,
dst_crs=mgrs_crs,
resampling=Resampling.cubic)
mgrs_dem = DEM(dst_dem[0, :, :].astype(np.int16),
dst_dem[1, :, :].astype(np.int16),
dst_dem[2, :, :].astype(np.int16), mgrs_crs,
dst_dem_transform)
return mgrs_dem
import numpy as np
import rasterio as rio
import src.sensorsio.srtm as srtm
from src.sensorsio import sentinel2
def test_srtm_id_to_name():
assert srtm.srtm_id_to_name(srtm.SRTMTileId(1, -2)) == "S02E001"
assert srtm.srtm_id_to_name(srtm.SRTMTileId(-1, -2)) == "S02W001"
assert srtm.srtm_id_to_name(srtm.SRTMTileId(-1, 12)) == "N12W001"
assert srtm.SRTMTileId(1, -2).name() == "S02E001"
assert srtm.SRTMTileId(-1, -2).name() == "S02W001"
assert srtm.SRTMTileId(-1, 12).name() == "N12W001"
def test_crs_from_mgrs():
assert srtm.crs_for_mgrs_tile('31TDH').to_authority() == ('EPSG', '32631')
assert srtm.get_crs_mgrs_tile('31TDH').to_authority() == ('EPSG', '32631')
def test_mgrs_transform():
assert srtm.mgrs_transform('31TDH') == rio.Affine(10.0, 0.0, 399960.0, 0.0,
-10.0, 4800000.0)
assert srtm.get_transform_mgrs_tile('31TDH') == rio.Affine(
10.0, 0.0, 399960.0, 0.0, -10.0, 4800000.0)
def test_srtm_tiles_from_mgrs_tile():
def build_tile_list(tile):
return [
srtm.srtm_id_to_name(tid)
for tid in srtm.srtm_tiles_from_mgrs_tile(tile)
]
return [tid.name() for tid in srtm.get_srtm_tiles_for_mgrs_tile(tile)]
assert build_tile_list("31TCJ") == [
'N43E000', 'N43E001', 'N44E000', 'N44E001'
......@@ -45,7 +40,7 @@ def test_srtm_tiles_from_mgrs_tile():
def test_generate_dem():
dem_handler = srtm.SRTM()
dem = dem_handler.get_dem_for_mgrs_tile("31TDH")
dem = dem_handler.get_dem_mgrs_tile("31TDH")
with rio.open("/tmp/dem.tif",
'w',
driver='GTiff',
......@@ -62,78 +57,7 @@ def test_generate_dem():
ds.write(dem.aspect, 3)
def test_dem_on_s2_tile():
from rasterio.warp import Resampling, reproject
RESOLUTION = 10
s2_ds = sentinel2.Sentinel2(
'/datalake/S2-L2A-THEIA/31TDH/2019/05/31/'
'SENTINEL2B_20190531-105916-927_L2A_T31TDH_C_V2-2/')
xrds = s2_ds.read_as_xarray([sentinel2.Sentinel2.B12],
resolution=RESOLUTION)
dst_dem = np.zeros((3, xrds.sizes['x'], xrds.sizes['y']))
dem_handler = srtm.SRTM()
s2_dem = dem_handler.get_dem_for_mgrs_tile(s2_ds.tile)
dst_dem, dst_dem_transform = reproject(s2_dem.as_stack(),
destination=dst_dem,
src_transform=s2_dem.transform,
src_crs=s2_dem.crs,
dst_transform=s2_ds.transform,
dst_crs=s2_ds.crs,
resampling=Resampling.cubic)
s2_dem = srtm.DEM(dst_dem[0, :, :].astype(np.int16),
dst_dem[1, :, :].astype(np.int16),
dst_dem[2, :, :].astype(np.int16), s2_ds.crs,
dst_dem_transform)
with rio.open("/tmp/s2_dem.tif",
'w',
driver='GTiff',
height=s2_dem.elevation.shape[0],
width=s2_dem.elevation.shape[1],
count=3,
nodata=-32768.0,
dtype=s2_dem.elevation.dtype,
compress='lzw',
crs=s2_dem.crs,
transform=s2_dem.transform) as ds:
ds.write(s2_dem.elevation, 1)
ds.write(s2_dem.slope, 2)
ds.write(s2_dem.aspect, 3)
def test_dem_on_mgrs_tile():
from rasterio.warp import Resampling, reproject
TILE = '31TDH'
mgrs_transform = srtm.mgrs_transform(TILE)
mgrs_crs = srtm.crs_for_mgrs_tile(TILE)
dst_dem = np.zeros((3, 10980, 10980))
dem_handler = srtm.SRTM()
s2_dem = dem_handler.get_dem_for_mgrs_tile(TILE)
dst_dem, dst_dem_transform = reproject(s2_dem.as_stack(),
destination=dst_dem,
src_transform=s2_dem.transform,
src_crs=s2_dem.crs,
dst_transform=mgrs_transform,
dst_crs=mgrs_crs,
resampling=Resampling.cubic)
s2_dem = srtm.DEM(dst_dem[0, :, :].astype(np.int16),
dst_dem[1, :, :].astype(np.int16),
dst_dem[2, :, :].astype(np.int16), mgrs_crs,
dst_dem_transform)
with rio.open("/tmp/mgrs_dem.tif",
'w',
driver='GTiff',
height=s2_dem.elevation.shape[0],
width=s2_dem.elevation.shape[1],
count=3,
nodata=-32768.0,
dtype=s2_dem.elevation.dtype,
compress='lzw',
crs=s2_dem.crs,
transform=s2_dem.transform) as ds:
ds.write(s2_dem.elevation, 1)
ds.write(s2_dem.slope, 2)
ds.write(s2_dem.aspect, 3)
s2_dem = srtm.get_dem_mgrs_tile(TILE)
s2_dem.write(f"/tmp/dem_{TILE}.tif")
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment