Repository: microsoft/torchgeo

Issue:
RandomGeoSampler bias
I've discovered two instances of bias in our current `RandomGeoSampler` and `RandomBatchGeoSampler` implementations.

### Area bias

In our current implementations, we first select a tile uniformly at random, then choose a random patch from that tile. This means that small tiles are just as likely to be sampled as large tiles. For most datasets, this is not an issue, as all tiles are approximately the same size. However, for `IntersectionDatasets`, the area of each intersection can vary widely. Tiles that are barely large enough to return a single sample from will be sampled no less often than massive tiles. This is particularly problematic for `RandomBatchGeoSampler` which would sample an entire batch of images from that very small tile.

This is an issue inherent in our current implementation, but is relatively easy to fix. The solution would be to use a weighted random sampler where weights are derived from the area of each image. Therefore, large tiles will be more likely to be sampled from than small tiles.

### Latitude bias

In certain 2D projections like Mercator, polar regions have the same area as equatorial regions. However, this is not the case on the actual 3D Earth. If we sample at random, we end up oversampling the poles relative to the equator. This has serious consequences for model training including models that are biased towards polar regions or that overestimate/underestimate certain climate patterns.

This is an issue inherent to 2D projections of the Earth in general, not necessarily an issue in TorchGeo. We could still try to do something about this however. One solution would be to force the R-tree index to be in an equal-area projection like Albers. Another solution would be to force the index to be in Mercator and to use a weighted random sampler where the weight comes from (the square root of?) the latitude.


Discussion:
I think latitude bias is going to be difficult to correct. It depends on the CRS being used (equal angle CRSs are affected but not equal area) so we can't simply use a weighted random sampler without first checking against a list of known equal angle CRSs. It's also unclear what to do for non-equal angle and non-equal area CRSs.
Correcting area bias will require the new `bbox.area` attributes introduced in 0.3.0.

Resolving patch:
diff --git a/torchgeo/samplers/batch.py b/torchgeo/samplers/batch.py
--- a/torchgeo/samplers/batch.py
+++ b/torchgeo/samplers/batch.py
@@ -4,9 +4,9 @@
 """TorchGeo batch samplers."""
 
 import abc
-import random
 from typing import Iterator, List, Optional, Tuple, Union
 
+import torch
 from rtree.index import Index, Property
 from torch.utils.data import Sampler
 
@@ -104,13 +104,20 @@ def __init__(
         self.batch_size = batch_size
         self.length = length
         self.hits = []
+        areas = []
         for hit in self.index.intersection(tuple(self.roi), objects=True):
             bounds = BoundingBox(*hit.bounds)
             if (
-                bounds.maxx - bounds.minx > self.size[1]
-                and bounds.maxy - bounds.miny > self.size[0]
+                bounds.maxx - bounds.minx >= self.size[1]
+                and bounds.maxy - bounds.miny >= self.size[0]
             ):
                 self.hits.append(hit)
+                areas.append(bounds.area)
+
+        # torch.multinomial requires float probabilities > 0
+        self.areas = torch.tensor(areas, dtype=torch.float)
+        if torch.sum(self.areas) == 0:
+            self.areas += 1
 
     def __iter__(self) -> Iterator[List[BoundingBox]]:
         """Return the indices of a dataset.
@@ -119,8 +126,9 @@ def __iter__(self) -> Iterator[List[BoundingBox]]:
             batch of (minx, maxx, miny, maxy, mint, maxt) coordinates to index a dataset
         """
         for _ in range(len(self)):
-            # Choose a random tile
-            hit = random.choice(self.hits)
+            # Choose a random tile, weighted by area
+            idx = torch.multinomial(self.areas, 1)
+            hit = self.hits[idx]
             bounds = BoundingBox(*hit.bounds)
 
             # Choose random indices within that tile
diff --git a/torchgeo/samplers/single.py b/torchgeo/samplers/single.py
--- a/torchgeo/samplers/single.py
+++ b/torchgeo/samplers/single.py
@@ -4,7 +4,6 @@
 """TorchGeo samplers."""
 
 import abc
-import random
 from typing import Callable, Iterable, Iterator, Optional, Tuple, Union
 
 import torch
@@ -105,13 +104,20 @@ def __init__(
 
         self.length = length
         self.hits = []
+        areas = []
         for hit in self.index.intersection(tuple(self.roi), objects=True):
             bounds = BoundingBox(*hit.bounds)
             if (
-                bounds.maxx - bounds.minx > self.size[1]
-                and bounds.maxy - bounds.miny > self.size[0]
+                bounds.maxx - bounds.minx >= self.size[1]
+                and bounds.maxy - bounds.miny >= self.size[0]
             ):
                 self.hits.append(hit)
+                areas.append(bounds.area)
+
+        # torch.multinomial requires float probabilities > 0
+        self.areas = torch.tensor(areas, dtype=torch.float)
+        if torch.sum(self.areas) == 0:
+            self.areas += 1
 
     def __iter__(self) -> Iterator[BoundingBox]:
         """Return the index of a dataset.
@@ -120,8 +126,9 @@ def __iter__(self) -> Iterator[BoundingBox]:
             (minx, maxx, miny, maxy, mint, maxt) coordinates to index a dataset
         """
         for _ in range(len(self)):
-            # Choose a random tile
-            hit = random.choice(self.hits)
+            # Choose a random tile, weighted by area
+            idx = torch.multinomial(self.areas, 1)
+            hit = self.hits[idx]
             bounds = BoundingBox(*hit.bounds)
 
             # Choose a random index within that tile
diff --git a/torchgeo/samplers/utils.py b/torchgeo/samplers/utils.py
--- a/torchgeo/samplers/utils.py
+++ b/torchgeo/samplers/utils.py
@@ -3,9 +3,10 @@
 
 """Common sampler utilities."""
 
-import random
 from typing import Tuple, Union
 
+import torch
+
 from ..datasets import BoundingBox
 
 
@@ -46,11 +47,18 @@ def get_random_bounding_box(
     t_size = _to_tuple(size)
 
     width = (bounds.maxx - bounds.minx - t_size[1]) // res
-    minx = random.randrange(int(width)) * res + bounds.minx
-    maxx = minx + t_size[1]
-
     height = (bounds.maxy - bounds.miny - t_size[0]) // res
-    miny = random.randrange(int(height)) * res + bounds.miny
+
+    minx = bounds.minx
+    miny = bounds.miny
+
+    # random.randrange crashes for inputs <= 0
+    if width > 0:
+        minx += torch.rand(1).item() * width * res
+    if height > 0:
+        miny += torch.rand(1).item() * height * res
+
+    maxx = minx + t_size[1]
     maxy = miny + t_size[0]
 
     mint = bounds.mint


Test patch:
diff --git a/tests/samplers/test_batch.py b/tests/samplers/test_batch.py
--- a/tests/samplers/test_batch.py
+++ b/tests/samplers/test_batch.py
@@ -127,6 +127,23 @@ def test_small_area(self) -> None:
         for _ in sampler:
             continue
 
+    def test_point_data(self) -> None:
+        ds = CustomGeoDataset()
+        ds.index.insert(0, (0, 0, 0, 0, 0, 0))
+        ds.index.insert(1, (1, 1, 1, 1, 1, 1))
+        sampler = RandomBatchGeoSampler(ds, 0, 2, 10)
+        for _ in sampler:
+            continue
+
+    def test_weighted_sampling(self) -> None:
+        ds = CustomGeoDataset()
+        ds.index.insert(0, (0, 0, 0, 0, 0, 0))
+        ds.index.insert(1, (0, 10, 0, 10, 0, 10))
+        sampler = RandomBatchGeoSampler(ds, 1, 2, 10)
+        for batch in sampler:
+            for bbox in batch:
+                assert bbox == BoundingBox(0, 10, 0, 10, 0, 10)
+
     @pytest.mark.slow
     @pytest.mark.parametrize("num_workers", [0, 1, 2])
     def test_dataloader(
diff --git a/tests/samplers/test_single.py b/tests/samplers/test_single.py
--- a/tests/samplers/test_single.py
+++ b/tests/samplers/test_single.py
@@ -122,6 +122,22 @@ def test_small_area(self) -> None:
         for _ in sampler:
             continue
 
+    def test_point_data(self) -> None:
+        ds = CustomGeoDataset()
+        ds.index.insert(0, (0, 0, 0, 0, 0, 0))
+        ds.index.insert(1, (1, 1, 1, 1, 1, 1))
+        sampler = RandomGeoSampler(ds, 0, 10)
+        for _ in sampler:
+            continue
+
+    def test_weighted_sampling(self) -> None:
+        ds = CustomGeoDataset()
+        ds.index.insert(0, (0, 0, 0, 0, 0, 0))
+        ds.index.insert(1, (0, 10, 0, 10, 0, 10))
+        sampler = RandomGeoSampler(ds, 1, 10)
+        for bbox in sampler:
+            assert bbox == BoundingBox(0, 10, 0, 10, 0, 10)
+
     @pytest.mark.slow
     @pytest.mark.parametrize("num_workers", [0, 1, 2])
     def test_dataloader(