suggested changes

microsoft · Dec 30, 2022 · a5a007c · a5a007c
1 parent dd30f0a
commit a5a007c
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 4 deletions.
diff --git a/torchgeo/datamodules/potsdam.py b/torchgeo/datamodules/potsdam.py
@@ -13,10 +13,9 @@
 from torch.utils.data import DataLoader, Dataset
 from torchvision.transforms import Compose
 
-from torchgeo.datasets.utils import collate_patches_per_tile
-from torchgeo.samplers.utils import _to_tuple
-
 from ..datasets import Potsdam2D
+from ..datasets.utils import collate_patches_per_tile
+from ..samplers.utils import _to_tuple
 from .utils import dataset_split
 
 

diff --git a/torchgeo/datasets/__init__.py b/torchgeo/datasets/__init__.py
@@ -99,6 +99,7 @@
 from .usavars import USAVars
 from .utils import (
     BoundingBox,
+    collate_patches_per_tile,
     concat_samples,
     merge_samples,
     stack_samples,
@@ -210,4 +211,5 @@
     "merge_samples",
     "stack_samples",
     "unbind_samples",
+    "collate_patches_per_tile",
 )
diff --git a/torchgeo/datasets/utils.py b/torchgeo/datasets/utils.py
@@ -229,7 +229,8 @@ def collate_patches_per_tile(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
         'train_batch_size' * 'num_patches_per_tile'
     """
     r_batch: Dict[str, Any] = default_collate(batch)  # type: ignore[no-untyped-call]
-    r_batch["image"] = rearrange(r_batch["image"], "b t c h w -> (b t) c h w")
+    if "mask" in r_batch:
+        r_batch["image"] = rearrange(r_batch["image"], "b t c h w -> (b t) c h w")
     r_batch["mask"] = rearrange(r_batch["mask"], "b t h w -> (b t) h w")
     return r_batch