Add MMEarth dataset

[nilsleh]

This PR adds the MMEarth dataset. Implementation is tested with the 100k version under the assumption that the dataset format is identical for the other versions.

There are some expected changes to the dataset so still a draft. But opening a PR already for some discussion points:

• As a multi modal dataset there are various configuration and selection schemes. The dataset has arguments for modality and band selection within modalities (the current implementation is a suggestion, but open for cleaner suggestions)
• For that same reason, we don't have a classic "image" and "label" sample dictionary because there are various ways of using this dataset across classification, segmentation and pixelwise regression. Is it reasonable to expect people to write transform functions that would define the sample as needed for their models?
• Normalization statistics are provided with the dataset. Normally we do standardization in the DataModule, however, given the complexity of various modalities and also two normalization schemes (z-score and min-max) I think it makes the dataset more useful when supplying this functionality in the dataset already
• The most complex part is handling the different modalities because some of them require some extra handling as seen here
• Additionally, need to decide what to do with the nan-values and clearly comment that
• currently this is just an implementation to support a non-geo dataset with standard index sampling, for other indexing schemes we can write an inherited class that also has a script that generates the needed parquet file and download it from HF
• probably things I am missing so feel free to add @adamjstewart @ando-shah

[nilsleh]

@adamjstewart and @ando-shah I think as a basic Dataset this does the job. For any more specific requirements we can subclass this to support other model specific needs.

[adamjstewart]

I added a new test to ruff

[adamjstewart]

Planning to get a working draft of SatlasPretrain first and then we can modify this PR to ensure it matches and meets @ando-shah's needs. These will be some of our first multi-modal datasets, so want to make sure we're doing things consistently (at least at an external user-facing API level).

[adamjstewart]

Validation is good, but not if it makes the code too complex.

[nilsleh]

Ando requested that Sentinel 1 should be separated into ascending and descending. That turned out to be a bit tricky, since MMEarth only stores "sentinel1" as a concatenation of ascending and descending, but I think I implemented it correctly to support band selection and normalization etc

[nilsleh]

As another point regarding Sentinel1:

• Sentinel 1 ascending and descending are stored separately, however, in the hdf5 file, there is just a single sentinel1 key that always has a shape of [8, 128, 128] so if one of ascending or descending or a subset of bands within the two is not available it is still stored as [1, 128,128] of nans in the dataset
• in order to not return Nans from the __getitem__ , only non-nan data is now returned, i.e. if I request {'sentinel1_asc': ['VV', 'VH', 'HH', 'HV']} in the dataset init but for the particular tile only {'sentinel1_asc': ['VV', 'VH']} is available, then a [2, 128, 128] tensor will be returned
• additionally, the sample contains a dict with the sensor and bands actually found in this particular sample

[adamjstewart]

Data loader seems much more complicated than it needs to be. Does it satisfy all of Ando's requirements?

[nilsleh]

Data loader seems much more complicated than it needs to be. Does it satisfy all of Ando's requirements?

Yeah all the feedback I have gotten from him, I have incorporated. I agree it is complicated but it's also loading data from >10 modalities so not a typical single image/target dataset. He said he will try it now and give any other feedback.

[nilsleh]

@ando-shah found a bug in my implementation. The available band names were different from the sample specific band name, for example in era5, they have the current or previous date of the sample included in the band name, so not a general name, and therefore it was not correctly retrieved.

[adamjstewart]

At the moment, all_modality_bands isn't documented. Should it be?

[nilsleh]

I meant to point to the class variable all_modality_bands

[adamjstewart]

Yes, but all_modality_bands isn't documented: https://torchgeo--2202.org.readthedocs.build/en/2202/api/datasets.html#torchgeo.datasets.MMEarth

I know you and I would just look at the source code instead of the docs, but most users will only ever look at the docs, and currently this references a variable that isn't documented.

If you want to document it, you can use something like:

#: List of valid modality bands
all_modality_bands = ...

Then you could link to it like:

modality_bands: dictionary of modality bands, see :attr:`all_modality_bands`

[nilsleh]

@adamjstewart @ando-shah while trying to write a subclass based on a metadata file, an annoying part is the image_{} and mask_{} naming for the sample keys, because you need to remove them again to match the meta data naming of the modalities. So since there isn't a torchgeo trainer that works out of the box on this dataset anyways, maybe we can just return the modality name itself? Just wondering if it potentially creates more confusion to have a chosen set of prefixes for the modalities, when some could be either etc.

[adamjstewart]

maybe we can just return the modality name itself?

Then kornia.augmentation.AugmentationSequential will no longer work.