Repository: microsoft/torchgeo

Issue:
Add support for float type raster datasets
I'm trying to load a custom DSM raster dataset of type float32.

I'm implementing a new class based on RasterDataset like so:
```
from torchgeo.datasets import RasterDataset

class DsmData(RasterDataset):
    filename_glob = "*.tif"
```

When I call `__getitem__()` on a DsmData object, I get an array containing int32 data instead.
I traced this down to the following line in the RasterDataset definition
https://github.com/microsoft/torchgeo/blob/3f7e525fbd01dddd25804e7a1b7634269ead1760/torchgeo/datasets/geo.py#L464

All raster datasets are forced into the int32 type which should not happen.

Having RasterDatasets with different types will probably cause problems when applying union or intersection operations.
But since a custom `collate_fn` can be defined, the user is able to provide a solution for this.

I achieved the desired behavior simply by removing the mentioned line.


Discussion:
Oof, good catch. We should definitely change this. Do you want to open a PR to get this started?
> Having RasterDatasets with different types will probably cause problems when applying union or intersection operations.

Nope, this won't be a problem. All of our collation functions use torch operators that can handle various dtypes:
```console
$ python
>>> import torch
>>> a = torch.zeros(3, 3, dtype=torch.int32)
>>> b = torch.zeros(3, 3, dtype=torch.float32)
>>> torch.stack((a, b))
tensor([[[0., 0., 0.],
         [0., 0., 0.],
         [0., 0., 0.]],

        [[0., 0., 0.],
         [0., 0., 0.],
         [0., 0., 0.]]])
>>> torch.maximum(a, b)
tensor([[0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.]])
>>> torch.cat((a, b))
tensor([[0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.]])
```

Resolving patch:
diff --git a/torchgeo/datasets/geo.py b/torchgeo/datasets/geo.py
--- a/torchgeo/datasets/geo.py
+++ b/torchgeo/datasets/geo.py
@@ -461,7 +461,12 @@ def _merge_files(self, filepaths: Sequence[str], query: BoundingBox) -> Tensor:
             )
         else:
             dest, _ = rasterio.merge.merge(vrt_fhs, bounds, self.res)
-        dest = dest.astype(np.int32)
+
+        # fix numpy dtypes which are not supported by pytorch tensors
+        if dest.dtype == np.uint16:
+            dest = dest.astype(np.int32)
+        elif dest.dtype == np.uint32:
+            dest = dest.astype(np.int64)
 
         tensor: Tensor = torch.tensor(dest)  # type: ignore[attr-defined]
         return tensor


Test patch:
diff --git a/tests/data/raster/data.py b/tests/data/raster/data.py
new file mode 100644
--- /dev/null
+++ b/tests/data/raster/data.py
@@ -0,0 +1,42 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+
+import numpy as np
+import rasterio
+import rasterio.transform
+from torchvision.datasets.utils import calculate_md5
+
+
+def generate_test_data(fn: str) -> str:
+    """Creates test data with uint32 datatype.
+
+    Args:
+        fn (str): Filename to write
+
+    Returns:
+        str: md5 hash of created archive
+    """
+    profile = {
+        "driver": "GTiff",
+        "dtype": "uint32",
+        "count": 1,
+        "crs": "epsg:4326",
+        "transform": rasterio.transform.from_bounds(0, 0, 1, 1, 1, 1),
+        "height": 4,
+        "width": 4,
+        "compress": "lzw",
+        "predictor": 2,
+    }
+
+    with rasterio.open(fn, "w", **profile) as f:
+        f.write(np.random.randint(0, 2**32 - 1, size=(1, 4, 4)))
+
+    md5: str = calculate_md5(fn)
+    return md5
+
+
+if __name__ == "__main__":
+    md5_hash = generate_test_data(os.path.join(os.getcwd(), "test0.tif"))
+    print(md5_hash)
diff --git a/tests/data/raster/test0.tif b/tests/data/raster/test0.tif
new file mode 100644
Binary files /dev/null and b/tests/data/raster/test0.tif differ
diff --git a/tests/datasets/test_geo.py b/tests/datasets/test_geo.py
--- a/tests/datasets/test_geo.py
+++ b/tests/datasets/test_geo.py
@@ -159,6 +159,11 @@ def sentinel(self, request: SubRequest) -> Sentinel2:
         cache = request.param
         return Sentinel2(root, bands=bands, transforms=transforms, cache=cache)
 
+    @pytest.fixture()
+    def custom_dtype_ds(self) -> RasterDataset:
+        root = os.path.join("tests", "data", "raster")
+        return RasterDataset(root)
+
     def test_getitem_single_file(self, naip: NAIP) -> None:
         x = naip[naip.bounds]
         assert isinstance(x, dict)
@@ -171,6 +176,12 @@ def test_getitem_separate_files(self, sentinel: Sentinel2) -> None:
         assert isinstance(x["crs"], CRS)
         assert isinstance(x["image"], torch.Tensor)
 
+    def test_getitem_uint_dtype(self, custom_dtype_ds: RasterDataset) -> None:
+        x = custom_dtype_ds[custom_dtype_ds.bounds]
+        assert isinstance(x, dict)
+        assert isinstance(x["image"], torch.Tensor)
+        assert x["image"].dtype == torch.int64  # type: ignore[attr-defined]
+
     def test_invalid_query(self, sentinel: Sentinel2) -> None:
         query = BoundingBox(0, 0, 0, 0, 0, 0)
         with pytest.raises(