Skip to content

Commit

Permalink
Merge pull request #21 from TUW-GEO/vrt_stack_reading_fixes
Browse files Browse the repository at this point in the history
VRT stack reading fixes
  • Loading branch information
cnavacch authored Jul 3, 2023
2 parents e25384f + b490694 commit 286812f
Show file tree
Hide file tree
Showing 3 changed files with 33 additions and 15 deletions.
6 changes: 6 additions & 0 deletions CHANGELOG.rst
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,12 @@
Changelog
=========

Version 1.1.0
=============

- enabled reading GeoTIFF files from ZIP files
- fixed bug in parallel VRT stack reading (#16) and layer assignment (#19)

Version 1.0.1
=============

Expand Down
7 changes: 6 additions & 1 deletion src/veranda/raster/mosaic/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -136,10 +136,15 @@ def data_geom(self) -> RasterGeometry:
""" Raster/tile geometry of the raster mosaic files. """
return self._data_geom

@property
def layer_ids(self) -> list:
""" List : Sorted layers. """
return list(sorted(self._file_register[self._file_dim].unique()))

@property
def n_layers(self) -> int:
""" Maximum number of layers. """
return self._file_register.groupby([self._tile_dim])[self._tile_dim].count().max()
return len(self.layer_ids)

@property
def n_tiles(self) -> int:
Expand Down
35 changes: 21 additions & 14 deletions src/veranda/raster/mosaic/geotiff.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,13 +24,14 @@
PROC_OBJS = {}


def read_init(fr, am, sm, sd, td, ad, dc, dk):
def read_init(fr, am, sm, sd, td, si, ad, dc, dk):
""" Helper method for setting the entries of global variable `PROC_OBJS` to be available during multiprocessing. """
PROC_OBJS['global_file_register'] = fr
PROC_OBJS['access_map'] = am
PROC_OBJS['shm_map'] = sm
PROC_OBJS['stack_dimension'] = sd
PROC_OBJS['tile_dimension'] = td
PROC_OBJS['stack_ids'] = si
PROC_OBJS['auto_decode'] = ad
PROC_OBJS['decoder'] = dc
PROC_OBJS['decoder_kwargs'] = dk
Expand Down Expand Up @@ -392,7 +393,7 @@ def __read_vrt_stack(self, access_map, shm_map, n_cores=1,
global_file_register = self._file_register

with Pool(n_cores, initializer=read_init, initargs=(global_file_register, access_map, shm_map,
self._file_dim, self._tile_dim,
self._file_dim, self._tile_dim, self.layer_ids,
auto_decode, decoder, decoder_kwargs)) as p:
p.map(read_vrt_stack, access_map.keys())

Expand Down Expand Up @@ -655,26 +656,32 @@ def read_vrt_stack(tile_id):
access_map = PROC_OBJS['access_map']
shm_map = PROC_OBJS['shm_map']
tile_dimension = PROC_OBJS['tile_dimension']
stack_dimension = PROC_OBJS['stack_dimension']
stack_ids = PROC_OBJS['stack_ids']

gt_access = access_map[tile_id]
bands = list(shm_map.keys())
file_register = global_file_register.loc[global_file_register[tile_dimension] == tile_id]

path = tempfile.gettempdir()
vrt_filepath = os.path.join(path, f"{datetime.utcnow().strftime('%Y%m%d%H%M%S')}-{uuid.uuid4().hex}.vrt")
while os.path.exists(vrt_filepath):
if len(file_register) > 0:
path = tempfile.gettempdir()
vrt_filepath = os.path.join(path, f"{datetime.utcnow().strftime('%Y%m%d%H%M%S')}-{uuid.uuid4().hex}.vrt")
filepaths = list(file_register['filepath'])
create_vrt_file(filepaths, vrt_filepath, gt_access.src_shape, gt_access.src_wkt, gt_access.src_geotrans,
bands=bands)
while os.path.exists(vrt_filepath):
vrt_filepath = os.path.join(path, f"{datetime.utcnow().strftime('%Y%m%d%H%M%S')}-{uuid.uuid4().hex}.vrt")

filepaths = list(file_register['filepath'])
create_vrt_file(filepaths, vrt_filepath, gt_access.src_shape, gt_access.src_wkt, gt_access.src_geotrans,
bands=bands)

src = gdal.Open(vrt_filepath, gdal.GA_ReadOnly)
vrt_data = src.ReadAsArray(*gt_access.gdal_args)
for band in bands:
_assign_vrt_stack_per_band(tile_id, band, src, vrt_data)
layer_ids = [stack_ids.index(stack_id) for stack_id in file_register[stack_dimension]]

src = gdal.Open(vrt_filepath, gdal.GA_ReadOnly)
vrt_data = src.ReadAsArray(*gt_access.gdal_args)
for band in bands:
_assign_vrt_stack_per_band(tile_id, layer_ids, band, src, vrt_data)


def _assign_vrt_stack_per_band(tile_id, band, src, vrt_data):
def _assign_vrt_stack_per_band(tile_id, layer_ids, band, src, vrt_data):
"""
Assigns loaded raster data to shared memory array for a specific band.
Expand Down Expand Up @@ -718,7 +725,7 @@ def _assign_vrt_stack_per_band(tile_id, band, src, vrt_data):

shm_rar, shm_ar_shape = shm_map[band]
shm_data = np.frombuffer(shm_rar, dtype=dtype).reshape(shm_ar_shape)
shm_data[:, gt_access.dst_row_slice, gt_access.dst_col_slice] = band_data
shm_data[layer_ids, gt_access.dst_row_slice, gt_access.dst_col_slice] = band_data


def read_single_files(file_idx):
Expand Down

0 comments on commit 286812f

Please sign in to comment.