Add dtype field to RasterDataset

tests/datasets/test_geo.py

@@ -227,6 +227,11 @@ def test_no_all_bands(self) -> None:
         with pytest.raises(AssertionError, match=msg):
             CustomSentinelDataset(root, bands=bands, transforms=transforms, cache=cache)
 
+    def test_dtype_warning(self, custom_dtype_ds: RasterDataset) -> None:
+        custom_dtype_ds.dtype = torch.int32
+        with pytest.warns(UserWarning, match="Custom dtype is explicitly set*"):
+            custom_dtype_ds[custom_dtype_ds.bounds]
+
 
 class TestVectorDataset:
     @pytest.fixture(scope="class")

torchgeo/datasets/chesapeake.py

@@ -40,6 +40,7 @@ class Chesapeake(RasterDataset, abc.ABC):
     """
 
     is_image = False
+    dtype = torch.long
 
     # subclasses use the 13 class cmap by default
     cmap = {

torchgeo/datasets/geo.py

@@ -11,6 +11,7 @@
 import sys
 from collections.abc import Sequence
 from typing import Any, Callable, Optional, cast
+from warnings import warn
 
 import fiona
 import fiona.transform
@@ -296,6 +297,9 @@ class RasterDataset(GeoDataset):
     #: Color map for the dataset, used for plotting
     cmap: dict[int, tuple[int, int, int, int]] = {}
 
+    #: dtype to force onto the dataset (overrides the dtype of the file via a cast)
+    dtype: Optional[torch.dtype] = None
+
     def __init__(
         self,
         root: str = "data",
@@ -429,10 +433,20 @@ def __getitem__(self, query: BoundingBox) -> dict[str, Any]:
             data = self._merge_files(filepaths, query, self.band_indexes)
 
         sample = {"crs": self.crs, "bbox": query}
+
+        if self.dtype is not None:
+            if self.is_image:
+                warn(
+                    "Custom dtype is explicitly set, but dtype is only valid for mask"
+                    + " RasterDatasets."
+                )
+            else:
+                data = data.to(self.dtype)
+
         if self.is_image:
             sample["image"] = data.float()
         else:
-            sample["mask"] = data.long()
+            sample["mask"] = data
 
         if self.transforms is not None:
             sample = self.transforms(sample)