Implement QgsRasterLayer.as_numpy

qgis · Aug 20, 2024 · 996a2b2 · 996a2b2
1 parent 217210a
commit 996a2b2
Show file tree

Hide file tree

Showing 3 changed files with 97 additions and 6 deletions.
diff --git 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
@@ -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
@@ -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()