Implement QgsRasterLayer.as_numpy

This commit is contained in:
merydian 2024-08-20 11:53:12 +02:00 committed by Nyall Dawson
parent 9c794474f8
commit 21a276fff2
3 changed files with 91 additions and 0 deletions

View File

@ -515,12 +515,30 @@ try:
QgsRasterBlock.as_numpy = _raster_block_as_numpy
def _raster_layer_as_numpy(self, use_masking=True, bands: _typing.Optional[_typing.List[int]] = None):
arrays = []
band_range = bands if bands else range(self.bandCount())
for band in band_range:
block = self.dataProvider().block(band + 1, self.extent(), self.width(), self.height())
src_array = block.as_numpy(use_masking=use_masking)
arrays.append(src_array)
return _numpy.array(arrays) # This converts any maskedArrays to numpy.array
QgsRasterLayer.as_numpy = _raster_layer_as_numpy
except ModuleNotFoundError:
def _raster_block_as_numpy(self, use_masking:bool = True):
raise QgsNotSupportedException('QgsRasterBlock.as_numpy is not available, numpy is not installed on the system')
QgsRasterBlock.as_numpy = _raster_block_as_numpy
def _raster_layer_as_numpy(self, use_masking:bool = True, bands: _typing.Optional[_typing.List[int]] = None):
raise QgsNotSupportedException('QgsRasterLayer.as_numpy is not available, numpy is not installed on the system')
QgsRasterLayer.as_numpy = _raster_layer_as_numpy
QgsRasterBlock.as_numpy.__doc__ = """
Returns the block data as a numpy array.
@ -530,3 +548,12 @@ If `use_masking` is `True` then the returned array will be a numpy masked array,
.. versionadded:: 3.40
"""
QgsRasterLayer.as_numpy.__doc__ = """
Returns the layer data as a numpy array.
If `use_masking` is `True` then the returned arrays will be numpy masked arrays, masking the raster block's nodata values.
:raises QgsNotSupportedException: if numpy is not available on the system
.. versionadded:: 3.40
"""

View File

@ -522,12 +522,30 @@ try:
QgsRasterBlock.as_numpy = _raster_block_as_numpy
def _raster_layer_as_numpy(self, use_masking=True, bands: _typing.Optional[_typing.List[int]] = None):
arrays = []
band_range = bands if bands else range(self.bandCount())
for band in band_range:
block = self.dataProvider().block(band + 1, self.extent(), self.width(), self.height())
src_array = block.as_numpy(use_masking=use_masking)
arrays.append(src_array)
return _numpy.array(arrays) # This converts any maskedArrays to numpy.array
QgsRasterLayer.as_numpy = _raster_layer_as_numpy
except ModuleNotFoundError:
def _raster_block_as_numpy(self, use_masking:bool = True):
raise QgsNotSupportedException('QgsRasterBlock.as_numpy is not available, numpy is not installed on the system')
QgsRasterBlock.as_numpy = _raster_block_as_numpy
def _raster_layer_as_numpy(self, use_masking:bool = True, bands: _typing.Optional[_typing.List[int]] = None):
raise QgsNotSupportedException('QgsRasterLayer.as_numpy is not available, numpy is not installed on the system')
QgsRasterLayer.as_numpy = _raster_layer_as_numpy
QgsRasterBlock.as_numpy.__doc__ = """
Returns the block data as a numpy array.
@ -537,3 +555,12 @@ If `use_masking` is `True` then the returned array will be a numpy masked array,
.. versionadded:: 3.40
"""
QgsRasterLayer.as_numpy.__doc__ = """
Returns the layer data as a numpy array.
If `use_masking` is `True` then the returned arrays will be numpy masked arrays, masking the raster block's nodata values.
:raises QgsNotSupportedException: if numpy is not available on the system
.. versionadded:: 3.40
"""

View File

@ -1264,6 +1264,43 @@ class TestQgsRasterLayerTransformContext(QgisTestCase):
for _ in range(num):
layer.readLayerXml(map_layer_element, context)
def test_as_numpy(self):
layer = QgsRasterLayer(self.rpath, 'raster')
arrays = layer.as_numpy()
self.assertEqual(type(arrays[5]), np.ndarray)
self.assertEqual(arrays.shape, (9, 200, 200))
self.assertEqual(arrays[0].dtype, np.int8)
# test with bands parameter
arrays = layer.as_numpy(bands=[1, 3])
self.assertEqual(type(arrays[0]), np.ndarray)
self.assertEqual(arrays.shape, (2, 200, 200))
self.assertEqual(arrays[0].dtype, np.int8)
path = os.path.join(unitTestDataPath('raster'),
'rgb_with_mask.tif')
layer = QgsRasterLayer(path, QFileInfo(path).baseName())
arrays = layer.as_numpy()
self.assertEqual(type(arrays[0]), np.ndarray)
self.assertEqual(arrays.shape, (4, 150, 162))
self.assertEqual(arrays[0].dtype, np.int8)
path = os.path.join(unitTestDataPath('raster'),
'rnd_percentile_raster5_float64.tif')
layer = QgsRasterLayer(path, QFileInfo(path).baseName())
arrays = layer.as_numpy()
self.assertEqual(type(arrays[0]), np.ndarray) # All maskedArrays are converted to numpy.array
self.assertEqual(arrays.shape, (1, 4, 4))
self.assertEqual(arrays[0].dtype, np.float64)
path = os.path.join(unitTestDataPath('raster'),
'rnd_percentile_raster5_float64.tif')
layer = QgsRasterLayer(path, QFileInfo(path).baseName())
arrays = layer.as_numpy(use_masking=False)
self.assertEqual(type(arrays[0]), np.ndarray)
self.assertEqual(arrays.shape, (1, 4, 4))
self.assertEqual(arrays[0].dtype, np.float64)
if __name__ == '__main__':
unittest.main()