Add pixel sampling mode

torchgeo/samplers/batch.py

@@ -12,6 +12,7 @@
 
 from torchgeo.datasets.geo import GeoDataset
 from torchgeo.datasets.utils import BoundingBox
+from torchgeo.samplers.constants import SIZE_IN_CRS_UNITS
 
 from .utils import _to_tuple, get_random_bounding_box
 
@@ -73,6 +74,7 @@ def __init__(
         batch_size: int,
         length: int,
         roi: Optional[BoundingBox] = None,
+        sample_mode: int = SIZE_IN_CRS_UNITS,
     ) -> None:
         """Initialize a new Sampler instance.
 
@@ -95,6 +97,7 @@ def __init__(
         self.size = _to_tuple(size)
         self.batch_size = batch_size
         self.length = length
+        self.sample_mode = sample_mode
         self.hits = list(self.index.intersection(tuple(self.roi), objects=True))
 
     def __iter__(self) -> Iterator[List[BoundingBox]]:
@@ -112,7 +115,9 @@ def __iter__(self) -> Iterator[List[BoundingBox]]:
             batch = []
             for _ in range(self.batch_size):
 
-                bounding_box = get_random_bounding_box(bounds, self.size, self.res)
+                bounding_box = get_random_bounding_box(
+                    bounds, self.size, self.res, self.sample_mode
+                )
                 batch.append(bounding_box)
 
             yield batch

torchgeo/samplers/constants.py

@@ -0,0 +1,6 @@
+from torch.utils.data import Sampler
+
+Sampler.__module__ = "torch.utils.data"
+
+SIZE_IN_PIXELS = 0
+SIZE_IN_CRS_UNITS = 1

torchgeo/samplers/single.py

@@ -13,6 +13,7 @@
 from torchgeo.datasets.geo import GeoDataset
 from torchgeo.datasets.utils import BoundingBox
 
+from .constants import SIZE_IN_CRS_UNITS
 from .utils import _to_tuple, get_random_bounding_box
 
 # https://github.com/pytorch/pytorch/issues/60979
@@ -75,6 +76,7 @@ def __init__(
         size: Union[Tuple[float, float], float],
         length: int,
         roi: Optional[BoundingBox] = None,
+        sample_mode: int = SIZE_IN_CRS_UNITS,
     ) -> None:
         """Initialize a new Sampler instance.
 
@@ -95,6 +97,7 @@ def __init__(
         super().__init__(dataset, roi)
         self.size = _to_tuple(size)
         self.length = length
+        self.sample_mode = sample_mode
         self.hits = list(self.index.intersection(tuple(self.roi), objects=True))
 
     def __iter__(self) -> Iterator[BoundingBox]:
@@ -109,7 +112,9 @@ def __iter__(self) -> Iterator[BoundingBox]:
             bounds = BoundingBox(*hit.bounds)
 
             # Choose a random index within that tile
-            bounding_box = get_random_bounding_box(bounds, self.size, self.res)
+            bounding_box = get_random_bounding_box(
+                bounds, self.size, self.res, self.sample_mode
+            )
 
             yield bounding_box
 

torchgeo/samplers/utils.py

@@ -8,6 +8,8 @@
 
 from torchgeo.datasets.utils import BoundingBox
 
+from .constants import SIZE_IN_PIXELS, SIZE_IN_CRS_UNITS
+
 
 def _to_tuple(value: Union[Tuple[float, float], float]) -> Tuple[float, float]:
     """Convert value to a tuple if it is not already a tuple.
@@ -25,7 +27,10 @@ def _to_tuple(value: Union[Tuple[float, float], float]) -> Tuple[float, float]:
 
 
 def get_random_bounding_box(
-    bounds: BoundingBox, size: Union[Tuple[float, float], float], res: float
+    bounds: BoundingBox,
+    size: Union[Tuple[float, float], float],
+    res: float,
+    sample_mode: int,
 ) -> BoundingBox:
     """Returns a random bounding box within a given bounding box.
 
@@ -39,6 +44,7 @@ def get_random_bounding_box(
     Args:
         bounds: the larger bounding box to sample from
         size: the size of the bounding box to sample
+        sample_mode: whether to sample in pixel space or CRS unit space
 
     Returns:
         randomly sampled bounding box from the extent of the input
@@ -47,11 +53,17 @@ def get_random_bounding_box(
 
     width = (bounds.maxx - bounds.minx - t_size[1]) // res
     minx = random.randrange(int(width)) * res + bounds.minx
-    maxx = minx + t_size[1]
+    if sample_mode == SIZE_IN_CRS_UNITS:
+        maxx = minx + t_size[1]
+    elif sample_mode == SIZE_IN_PIXELS:
+        maxx = minx + t_size[1] * res
 
     height = (bounds.maxy - bounds.miny - t_size[0]) // res
     miny = random.randrange(int(height)) * res + bounds.miny
-    maxy = miny + t_size[0]
+    if sample_mode == SIZE_IN_CRS_UNITS:
+        maxy = miny + t_size[0]
+    elif sample_mode == SIZE_IN_PIXELS:
+        maxy = miny + t_size[0] * res
 
     mint = bounds.mint
     maxt = bounds.maxt