From d27ee30a5f8839a6cf5c96175bea93d0c3716b6d Mon Sep 17 00:00:00 2001
From: "Adam J. Stewart" <ajstewart426@gmail.com>
Date: Wed, 10 Jul 2024 17:41:05 +0200
Subject: [PATCH] NumPy 2 support (#2151)

---
 pyproject.toml                   | 2 +-
 torchgeo/datasets/biomassters.py | 2 +-
 torchgeo/datasets/dfc2022.py     | 2 +-
 torchgeo/datasets/eurocrops.py   | 6 ++++--
 4 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 8a85c53a195..eb43ce1cbdd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -268,7 +268,7 @@ quote-style = "single"
 skip-magic-trailing-comma = true
 
 [tool.ruff.lint]
-extend-select = ["D", "I", "UP"]
+extend-select = ["D", "I", "NPY201", "UP"]
 
 [tool.ruff.lint.per-file-ignores]
 "docs/**" = ["D"]
diff --git a/torchgeo/datasets/biomassters.py b/torchgeo/datasets/biomassters.py
index bb975c8002b..ab440648a17 100644
--- a/torchgeo/datasets/biomassters.py
+++ b/torchgeo/datasets/biomassters.py
@@ -196,7 +196,7 @@ def _load_target(self, filename: str) -> Tensor:
             target mask
         """
         with rasterio.open(os.path.join(self.root, 'train_agbm', filename), 'r') as src:
-            arr: np.typing.NDArray[np.float_] = src.read()
+            arr: np.typing.NDArray[np.float64] = src.read()
 
         target = torch.from_numpy(arr).float()
         return target
diff --git a/torchgeo/datasets/dfc2022.py b/torchgeo/datasets/dfc2022.py
index 697ddeb0fb5..b9cd1556f9f 100644
--- a/torchgeo/datasets/dfc2022.py
+++ b/torchgeo/datasets/dfc2022.py
@@ -235,7 +235,7 @@ def _load_image(self, path: str, shape: Sequence[int] | None = None) -> Tensor:
             the image
         """
         with rasterio.open(path) as f:
-            array: np.typing.NDArray[np.float_] = f.read(
+            array: np.typing.NDArray[np.float64] = f.read(
                 out_shape=shape, out_dtype='float32', resampling=Resampling.bilinear
             )
             tensor = torch.from_numpy(array)
diff --git a/torchgeo/datasets/eurocrops.py b/torchgeo/datasets/eurocrops.py
index daa1987e3c8..0082dd152b9 100644
--- a/torchgeo/datasets/eurocrops.py
+++ b/torchgeo/datasets/eurocrops.py
@@ -243,10 +243,12 @@ def plot(
 
         fig, axs = plt.subplots(nrows=1, ncols=ncols, figsize=(4, 4))
 
-        def apply_cmap(arr: 'np.typing.NDArray[Any]') -> 'np.typing.NDArray[np.float_]':
+        def apply_cmap(
+            arr: 'np.typing.NDArray[Any]',
+        ) -> 'np.typing.NDArray[np.float64]':
             # Color 0 as black, while applying default color map for the class indices.
             cmap = plt.get_cmap('viridis')
-            im: np.typing.NDArray[np.float_] = cmap(arr / len(self.class_map))
+            im: np.typing.NDArray[np.float64] = cmap(arr / len(self.class_map))
             im[arr == 0] = 0
             return im