Sampling from UnionDataset fails - even when it shouldn't?

[PenguinJunk]

We want to create a DataLoader for multiple datasets but constantly get an IndexError thrown when sampling with the DataLoader over the UnionDataset: ("query: {query} not found in index with bounds: {self.bounds}"). This happens even when the query bounding box appears to fall within the bounds (according to the error message).

I am using our organisation's own data and working in AzureML and so need to use datasets registered in the AzureML workspace and mounting these (hence the first function).

Sampling with the DataLoader works fine if applied to the flight blocks individually. I have tested with datasets that are geographically adjacent. I wanted to test against the Landsat data given as an example in the docs, but finding the right data and downloading was too much of a challenge today! Hence I am not yet logging this as a issue.

Below is the code snippet and resulting error. I am also posted other examples of the error string. Any idea where I may be going wrong?

def mount_data(dataset, path, dataset_name):
    mount_loc = os.path.join(path, dataset_name)
    dataset_mnt = dataset.mount(mount_loc)
    dataset_mnt.start()
    dataset_mpt = dataset_mnt.mount_point
    return dataset_mnt, dataset_mpt

dataset1_name = 'ds1'
dataset2_name = 'ds2'

dataset1 = Dataset.get_by_name(workspace, name=dataset1_name)
dataset1_mnt, dataset1_mpt = mount_data(dataset1, '/tmp/', dataset1_name)
flightblock1 = RSSflightblock(root=dataset1_mpt)

dataset2 = Dataset.get_by_name(workspace, name=dataset2_name)
dataset2_mnt, dataset2_mpt = mount_data(dataset2, '/tmp/', dataset2_name)
flightblock2 = RSSflightblock(root=dataset2_mpt)

uniondataset = flightblock1 | flightblock2

sampler = RandomGeoSampler(uniondataset, size=256 , length=100)
dataloader = DataLoader(uniondataset, batch_size=16, sampler=sampler, collate_fn=stack_samples)

\# visual test
for sample in dataloader:
    image = torch.div(sample['image'],65536)
    fig, ax = plt.subplots(1,2)

IndexError: query: BoundingBox(minx=532473.1807103157, maxx=532537.1807103157, miny=159300.4143061638, maxy=159364.4143061638, mint=0.0, maxt=9.223372036854776e+18) not found in index with bounds: BoundingBox(minx=469000.0, maxx=534000.0, miny=149000.0, maxy=161000.0, mint=0.0, maxt=9.223372036854776e+18)
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Input In [24], in <cell line: 2>()
      1 # visual test
----> 2 for sample in dataloader:
      3     image = torch.div(sample['image'],65536)
      4     fig, ax = plt.subplots(1,2)

File /anaconda/envs/azureml_py38_PT_TF/lib/python3.8/site-packages/torch/utils/data/dataloader.py:530, in _BaseDataLoaderIter.__next__(self)
    528 if self._sampler_iter is None:
    529     self._reset()
--> 530 data = self._next_data()
    531 self._num_yielded += 1
    532 if self._dataset_kind == _DatasetKind.Iterable and \
    533         self._IterableDataset_len_called is not None and \
    534         self._num_yielded > self._IterableDataset_len_called:

File /anaconda/envs/azureml_py38_PT_TF/lib/python3.8/site-packages/torch/utils/data/dataloader.py:570, in _SingleProcessDataLoaderIter._next_data(self)
    568 def _next_data(self):
    569     index = self._next_index()  # may raise StopIteration
--> 570     data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
    571     if self._pin_memory:
    572         data = _utils.pin_memory.pin_memory(data)

File /anaconda/envs/azureml_py38_PT_TF/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py:49, in _MapDatasetFetcher.fetch(self, possibly_batched_index)
     47 def fetch(self, possibly_batched_index):
     48     if self.auto_collation:
---> 49         data = [self.dataset[idx] for idx in possibly_batched_index]
     50     else:
     51         data = self.dataset[possibly_batched_index]

File /anaconda/envs/azureml_py38_PT_TF/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py:49, in <listcomp>(.0)
     47 def fetch(self, possibly_batched_index):
     48     if self.auto_collation:
---> 49         data = [self.dataset[idx] for idx in possibly_batched_index]
     50     else:
     51         data = self.dataset[possibly_batched_index]

File /anaconda/envs/azureml_py38_PT_TF/lib/python3.8/site-packages/torchgeo/datasets/geo.py:986, in UnionDataset.__getitem__(self, query)
    984 for ds in self.datasets:
    985     if ds.index.intersection(tuple(query)):
--> 986         samples.append(ds[query])
    988 return self.collate_fn(samples)

File /anaconda/envs/azureml_py38_PT_TF/lib/python3.8/site-packages/torchgeo/datasets/geo.py:397, in RasterDataset.__getitem__(self, query)
    394 filepaths = [hit.object for hit in hits]
    396 if not filepaths:
--> 397     raise IndexError(
    398         f"query: {query} not found in index with bounds: {self.bounds}"
    399     )
    401 if self.separate_files:
    402     data_list: List[Tensor] = []

IndexError: query: BoundingBox(minx=532473.1807103157, maxx=532537.1807103157, miny=159300.4143061638, maxy=159364.4143061638, mint=0.0, maxt=9.223372036854776e+18) not found in index with bounds: BoundingBox(minx=469000.0, maxx=534000.0, miny=149000.0, maxy=161000.0, mint=0.0, maxt=9.223372036854776e+18)

Other error strings (formatted):

IndexError: query: BoundingBox(minx=478078.5790810585, maxx=478142.5790810585, miny=160747.37424325943, 
maxy=160811.37424325943, mint=0.0, maxt=9.223372036854776e+18) not found in index with bounds:
 BoundingBox(minx=469000.0, maxx=551000.0, miny=159000.0, maxy=181000.0, mint=0.0, maxt=9.223372036854776e+18)

IndexError: query: BoundingBox(minx=531457.1875019073, maxx=531465.1875019073, miny=160265.65970230103,
 maxy=160273.65970230103, mint=0.0, maxt=9.223372036854776e+18) not found in index with bounds: 
BoundingBox(minx=469000.0, maxx=551000.0, miny=159000.0, maxy=181000.0, mint=0.0, maxt=9.223372036854776e+18)

IndexError: query: BoundingBox(minx=480619.23027420044, maxx=480627.23027420044, miny=160568.82683372498,
 maxy=160576.82683372498, mint=0.0, maxt=9.223372036854776e+18) not found in index with bounds: 
BoundingBox(minx=469000.0, maxx=551000.0, miny=159000.0, maxy=181000.0, mint=0.0, maxt=9.223372036854776e+18)

[PenguinJunk]

Sorry, I should have also added that RSSflightblock is our own class of Geospatial Dataset.

[adamjstewart]

I will look into incorporating a regex into the filename_glob (hints welcome!).

So there are two things with very different purposes:

1. filename_glob used to search for files
2. filename_regex used to locate parts of the filename like timestamp or band number

The filename_glob can only include glob characters, it can't include any regular expressions. If your filename does not include a timestamp or band number, you don't need to use filename_regex at all.

Btw, I greatly expanded this tutorial: https://torchgeo.readthedocs.io/en/latest/tutorials/custom_raster_dataset.html. Hopefully this should clear up some confusion on defining the dataset. I still think your dataset is correct though, I think there must be a bug in TorchGeo. If you can figure out how to reproduce with public data, or even fake data, let me know.

[PenguinJunk]

Thank you!

I have managed to reproduce the problem with a smaller dataset and will send you an email to discuss how to make these data available to you.

However, I also want to note that I looked at switching the filter from filename_regex into the filename_glob in the dataset class definition i.e.

# old way removed 
class RSSflightblock(RasterDataset):
    filename_glob = "*.tif"
    filename_regex = r"""^(?!.*CMASK).*$"""

# correct? way
class RSSflightblock(RasterDataset):
    filename_glob = "*[!CMASK]*.tif" 
    # also tried both of these:
    #filename_glob = "*.tif"
    #filename_glob = "*[!C][!M][!A][!S][!K]*tif" 
    filename_regex = r"" # also tried leaving this line out

But this results in the following error (which may be something else entirely):

DatasetIOShapeError: Dataset indexes and destination buffer are mismatched
---------------------------------------------------------------------------
DatasetIOShapeError                       Traceback (most recent call last)
Input In [13], in <cell line: 1>()
----> 1 for batch in dataloader:
      2     sample = unbind_samples(batch)[0]
      3     #print(torch.amax(sample["image"]))

File /anaconda/envs/azureml_py38_PT_TF/lib/python3.8/site-packages/torch/utils/data/dataloader.py:530, in _BaseDataLoaderIter.__next__(self)
    528 if self._sampler_iter is None:
    529     self._reset()
--> 530 data = self._next_data()
    531 self._num_yielded += 1
    532 if self._dataset_kind == _DatasetKind.Iterable and \
    533         self._IterableDataset_len_called is not None and \
    534         self._num_yielded > self._IterableDataset_len_called:

File /anaconda/envs/azureml_py38_PT_TF/lib/python3.8/site-packages/torch/utils/data/dataloader.py:570, in _SingleProcessDataLoaderIter._next_data(self)
    568 def _next_data(self):
    569     index = self._next_index()  # may raise StopIteration
--> 570     data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
    571     if self._pin_memory:
    572         data = _utils.pin_memory.pin_memory(data)

File /anaconda/envs/azureml_py38_PT_TF/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py:49, in _MapDatasetFetcher.fetch(self, possibly_batched_index)
     47 def fetch(self, possibly_batched_index):
     48     if self.auto_collation:
---> 49         data = [self.dataset[idx] for idx in possibly_batched_index]
     50     else:
     51         data = self.dataset[possibly_batched_index]

File /anaconda/envs/azureml_py38_PT_TF/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py:49, in <listcomp>(.0)
     47 def fetch(self, possibly_batched_index):
     48     if self.auto_collation:
---> 49         data = [self.dataset[idx] for idx in possibly_batched_index]
     50     else:
     51         data = self.dataset[possibly_batched_index]

File /anaconda/envs/azureml_py38_PT_TF/lib/python3.8/site-packages/torchgeo/datasets/geo.py:424, in RasterDataset.__getitem__(self, query)
    422     data = torch.cat(data_list)
    423 else:
--> 424     data = self._merge_files(filepaths, query)
    426 key = "image" if self.is_image else "mask"
    427 sample = {key: data, "crs": self.crs, "bbox": query}

File /anaconda/envs/azureml_py38_PT_TF/lib/python3.8/site-packages/torchgeo/datasets/geo.py:459, in RasterDataset._merge_files(self, filepaths, query)
    455     dest = src.read(
    456         out_shape=out_shape, window=from_bounds(*bounds, src.transform)
    457     )
    458 else:
--> 459     dest, _ = rasterio.merge.merge(vrt_fhs, bounds, self.res)
    461 # fix numpy dtypes which are not supported by pytorch tensors
    462 if dest.dtype == np.uint16:

File /anaconda/envs/azureml_py38_PT_TF/lib/python3.8/site-packages/rasterio/merge.py:362, in merge(datasets, bounds, res, nodata, dtype, precision, indexes, output_count, resampling, method, target_aligned_pixels, dst_path, dst_kwds)
    356     temp_height, temp_width = (
    357         dst_window_rnd_off.height,
    358         dst_window_rnd_off.width,
    359     )
    360     temp_shape = (src_count, temp_height, temp_width)
--> 362     temp_src = src.read(
    363         out_shape=temp_shape,
    364         window=src_window_rnd_shp,
    365         boundless=False,
    366         masked=True,
    367         indexes=indexes,
    368         resampling=resampling,
    369     )
    371 # 5. Copy elements of temp into dest
    372 roff, coff = (
    373     max(0, dst_window_rnd_off.row_off),
    374     max(0, dst_window_rnd_off.col_off),
    375 )
    
File rasterio/_io.pyx:610, in rasterio._io.DatasetReaderBase.read()

File rasterio/_io.pyx:934, in rasterio._io.DatasetReaderBase._read()

File rasterio/_io.pyx:160, in rasterio._io.io_multi_band()

DatasetIOShapeError: Dataset indexes and destination buffer are mismatched

[adamjstewart]

Replied to email. I think the old glob/regex are correct, but I'll investigate once I have access to the data.

[adamjstewart]

This bug should be fixed by #786

[metazool]

Thank you!