Repository: microsoft/torchgeo

Issue:
Can't add transforms to UnionDataset or IntersectionDataset
### Summary

Both UnionDataset and IntersectionDataset don't have a way to add transforms neither call `self.transforms` in their `__get_item__` method. The transforms from the merged or intersected datasets are called in their own `__get_item__` methods, but there are situations where having it on the UnionDataset or IntersectionDataset is more convenient and useful.

### Rationale

In some situations having transforms on each merged or intersected dataset might be enough, but here's a couple of cases where this is needed:

* Using the IntersectionDataset to add a mask to an image. For the transforms where information information from both mask and image are needed (i.e. filter mask where image has no data).
* Using UnionDataset to merge multiple RasterDatasets (or IntersectionDatasets) and add only one transforms to the whole dataset instead of on each one.

### Implementation

One possible way would be to just add the call to self.transforms in UnionDataset and IntersectionDataset:

```
if self.transforms is not None:
     sample = self.transforms(sample)
```

And let people assign a transform after doing the intersection/union by doing `my_dataset.transforms = my_transforms`. Or adding an `add_transforms` method to IntersectionDataset and UnionDataset to make it more explicit.

### Alternatives

One alternative is to add a transforms parameter to IntersectionDataset and UnionDataset, and also call it on `__get_item__`. But this forces people to use the constructor directly instead of doing & or |.

Another option is to create a TransformableDataset that takes a dataset and transforms as parameters. Calls the underlying datasets `__get_item__` and then transforms in its `__get_item__`.

### Additional information

_No response_


Discussion:
I like this idea. Another reason this is important is for things like random flips/rotations that need to be applied uniformly to both images and target labels.

> One possible way would be to just add the call to self.transforms in UnionDataset and IntersectionDataset... And let people assign a transform after doing the intersection/union by doing `my_dataset.transforms = my_transforms`

> One alternative is to add a transforms parameter to IntersectionDataset and UnionDataset

I like both of these ideas, and I don't see any reason why we can't do both. Want to submit a PR?
Yes you are right, it makes sense to have both. Happy to submit a PR!

Resolving patch:
diff --git a/torchgeo/datasets/geo.py b/torchgeo/datasets/geo.py
--- a/torchgeo/datasets/geo.py
+++ b/torchgeo/datasets/geo.py
@@ -822,6 +822,7 @@ def __init__(
         collate_fn: Callable[
             [Sequence[Dict[str, Any]]], Dict[str, Any]
         ] = concat_samples,
+        transforms: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None,
     ) -> None:
         """Initialize a new Dataset instance.
 
@@ -829,11 +830,16 @@ def __init__(
             dataset1: the first dataset
             dataset2: the second dataset
             collate_fn: function used to collate samples
+            transforms: a function/transform that takes input sample and its target as
+                entry and returns a transformed version
 
         Raises:
             ValueError: if either dataset is not a :class:`GeoDataset`
+
+        .. versionadded:: 0.4
+            The *transforms* parameter.
         """
-        super().__init__()
+        super().__init__(transforms)
         self.datasets = [dataset1, dataset2]
         self.collate_fn = collate_fn
 
@@ -892,7 +898,12 @@ def __getitem__(self, query: BoundingBox) -> Dict[str, Any]:
         # All datasets are guaranteed to have a valid query
         samples = [ds[query] for ds in self.datasets]
 
-        return self.collate_fn(samples)
+        sample = self.collate_fn(samples)
+
+        if self.transforms is not None:
+            sample = self.transforms(sample)
+
+        return sample
 
     def __str__(self) -> str:
         """Return the informal string representation of the object.
@@ -934,6 +945,7 @@ def __init__(
         collate_fn: Callable[
             [Sequence[Dict[str, Any]]], Dict[str, Any]
         ] = merge_samples,
+        transforms: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None,
     ) -> None:
         """Initialize a new Dataset instance.
 
@@ -941,11 +953,16 @@ def __init__(
             dataset1: the first dataset
             dataset2: the second dataset
             collate_fn: function used to collate samples
+            transforms: a function/transform that takes input sample and its target as
+                entry and returns a transformed version
 
         Raises:
             ValueError: if either dataset is not a :class:`GeoDataset`
+
+        .. versionadded:: 0.4
+            The *transforms* parameter.
         """
-        super().__init__()
+        super().__init__(transforms)
         self.datasets = [dataset1, dataset2]
         self.collate_fn = collate_fn
 
@@ -1005,7 +1022,12 @@ def __getitem__(self, query: BoundingBox) -> Dict[str, Any]:
             if list(ds.index.intersection(tuple(query))):
                 samples.append(ds[query])
 
-        return self.collate_fn(samples)
+        sample = self.collate_fn(samples)
+
+        if self.transforms is not None:
+            sample = self.transforms(sample)
+
+        return sample
 
     def __str__(self) -> str:
         """Return the informal string representation of the object.


Test patch:
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
@@ -383,7 +383,8 @@ class TestIntersectionDataset:
     def dataset(self) -> IntersectionDataset:
         ds1 = CustomGeoDataset()
         ds2 = CustomGeoDataset()
-        return ds1 & ds2
+        transforms = nn.Identity()
+        return IntersectionDataset(ds1, ds2, transforms=transforms)
 
     def test_getitem(self, dataset: IntersectionDataset) -> None:
         query = BoundingBox(0, 1, 2, 3, 4, 5)
@@ -437,7 +438,8 @@ class TestUnionDataset:
     def dataset(self) -> UnionDataset:
         ds1 = CustomGeoDataset(bounds=BoundingBox(0, 1, 0, 1, 0, 1))
         ds2 = CustomGeoDataset(bounds=BoundingBox(2, 3, 2, 3, 2, 3))
-        return ds1 | ds2
+        transforms = nn.Identity()
+        return UnionDataset(ds1, ds2, transforms=transforms)
 
     def test_getitem(self, dataset: UnionDataset) -> None:
         query = BoundingBox(0, 1, 0, 1, 0, 1)