From 996a2b27f98806cfe3f22d1000e34f8652bfda8a Mon Sep 17 00:00:00 2001 From: merydian Date: Tue, 20 Aug 2024 11:53:12 +0200 Subject: [PATCH] Implement QgsRasterLayer.as_numpy --- python/PyQt6/core/__init__.py.in | 33 ++++++++++++++++++++-- python/core/__init__.py.in | 33 ++++++++++++++++++++-- tests/src/python/test_qgsrasterlayer.py | 37 +++++++++++++++++++++++++ 3 files changed, 97 insertions(+), 6 deletions(-) diff --git a/python/PyQt6/core/__init__.py.in b/python/PyQt6/core/__init__.py.in index 4e71a2028802..da6a32df2bd6 100644 --- a/python/PyQt6/core/__init__.py.in +++ b/python/PyQt6/core/__init__.py.in @@ -43,7 +43,7 @@ from .additions.runtimeprofiler import ScopedRuntimeProfileContextManager from .additions.validitycheck import check from .additions.ranges import datetime_range_repr, date_range_repr -from typing import Optional, Union +import typing as _typing # Injections into classes QgsFeature.__geo_interface__ = property(mapping_feature) @@ -454,7 +454,7 @@ QgsException.__doc__ = "Defines a QGIS exception class." try: import numpy as _numpy - def _qgis_data_type_to_numeric_data_type(dataType: Qgis.DataType) -> Optional[_numpy.dtype]: + def _qgis_data_type_to_numeric_data_type(dataType: Qgis.DataType) -> _typing.Optional[_numpy.dtype]: qgis_to_numpy_dtype_dict = { Qgis.DataType.UnknownDataType: None, Qgis.DataType.Byte: _numpy.byte, @@ -474,7 +474,7 @@ try: } return qgis_to_numpy_dtype_dict[dataType] - def _raster_block_as_numpy(self, use_masking:bool = True) -> Union[_numpy.ndarray, _numpy.ma.MaskedArray]: + def _raster_block_as_numpy(self, use_masking:bool = True) -> _typing.Union[_numpy.ndarray, _numpy.ma.MaskedArray]: raster_dtype = _qgis_data_type_to_numeric_data_type(self.dataType()) if not raster_dtype: raise ValueError(f"The raster block data type '{str(self.dataType())}' is not compatible with NumPy arrays.") @@ -488,12 +488,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. @@ -503,3 +521,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 +""" diff --git a/python/core/__init__.py.in b/python/core/__init__.py.in index 554314692d6d..57406a029be6 100644 --- a/python/core/__init__.py.in +++ b/python/core/__init__.py.in @@ -43,7 +43,7 @@ from .additions.runtimeprofiler import ScopedRuntimeProfileContextManager from .additions.validitycheck import check from .additions.ranges import datetime_range_repr, date_range_repr -from typing import Optional, Union +import typing as _typing # Injections into classes QgsFeature.__geo_interface__ = property(mapping_feature) @@ -460,7 +460,7 @@ QgsException.__doc__ = "Defines a QGIS exception class." try: import numpy as _numpy - def _qgis_data_type_to_numeric_data_type(dataType: Qgis.DataType) -> Optional[_numpy.dtype]: + def _qgis_data_type_to_numeric_data_type(dataType: Qgis.DataType) -> _typing.Optional[_numpy.dtype]: qgis_to_numpy_dtype_dict = { Qgis.DataType.UnknownDataType: None, Qgis.DataType.Byte: _numpy.byte, @@ -480,7 +480,7 @@ try: } return qgis_to_numpy_dtype_dict[dataType] - def _raster_block_as_numpy(self, use_masking:bool = True) -> Union[_numpy.ndarray, _numpy.ma.MaskedArray]: + def _raster_block_as_numpy(self, use_masking:bool = True) -> _typing.Union[_numpy.ndarray, _numpy.ma.MaskedArray]: raster_dtype = _qgis_data_type_to_numeric_data_type(self.dataType()) if not raster_dtype: raise ValueError(f"The raster block data type '{str(self.dataType())}' is not compatible with NumPy arrays.") @@ -494,12 +494,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. @@ -509,3 +527,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 +""" diff --git a/tests/src/python/test_qgsrasterlayer.py b/tests/src/python/test_qgsrasterlayer.py index 393671575348..51cdcaf7fbb9 100644 --- a/tests/src/python/test_qgsrasterlayer.py +++ b/tests/src/python/test_qgsrasterlayer.py @@ -1264,6 +1264,43 @@ def test_read_xml_crash(self): 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()