diff --git a/docs/source/datasets/dsprites.rst b/docs/source/datasets/dsprites.rst
index 5c10940c..1fc9444a 100644
--- a/docs/source/datasets/dsprites.rst
+++ b/docs/source/datasets/dsprites.rst
@@ -12,17 +12,46 @@ dSprites
Overview
--------
-The dSprites dataset is a synthetic benchmark designed for **disentangled and unsupervised representation learning**. It consists of procedurally generated **binary black-and-white images** of simple 2D shapes, rendered under controlled and fully known generative factors.
+The dSprites dataset is a synthetic benchmark designed for **disentangled and unsupervised representation learning**. It consists of procedurally generated 2D shapes rendered under controlled and fully known generative factors.
The dataset contains **all possible combinations** of six latent factors of variation, with each combination appearing exactly once. This complete Cartesian product structure makes dSprites a standard benchmark for evaluating disentanglement, factor predictability, and interpretability of learned representations.
-- **Total images**: 737,280
-- **Image resolution**: 64×64 (binary)
+In stable-datasets, dSprites is exposed via the ``DSprites`` class, with **four variants** selectable via ``config_name``:
+
+- **Total images**: 737,280 per variant
+- **Image resolution**: 64x64 (original) or 64x64x3 (color, noise, scream)
+
+Variants
+--------
+
+.. list-table::
+ :header-rows: 1
+ :widths: 15 15 55
+
+ * - Variant
+ - Image Mode
+ - Description
+ * - ``original``
+ - Grayscale
+ - Binary black-and-white images (default)
+ * - ``color``
+ - RGB
+ - Object rendered with a random RGB color on a black background
+ * - ``noise``
+ - RGB
+ - White object on a random-noise background
+ * - ``scream``
+ - RGB
+ - Object rendered by inverting pixels on a random Scream painting patch
+
+.. image:: teasers/dsprites_teaser.gif
+ :align: center
+ :width: 90%
Latent Factors of Variation
---------------------------
+---------------------------
-The dataset is generated from six independent latent factors:
+All variants share the same six independent latent factors:
.. list-table::
:header-rows: 1
@@ -42,7 +71,7 @@ The dataset is generated from six independent latent factors:
- Linearly spaced in [0.5, 1.0]
* - ``orientation``
- {0, ..., 39}
- - Uniform in [0, 2π] radians
+ - Uniform in [0, 2pi] radians
* - ``posX``
- {0, ..., 31}
- Normalized position in [0, 1]
@@ -50,12 +79,6 @@ The dataset is generated from six independent latent factors:
- {0, ..., 31}
- Normalized position in [0, 1]
-Each image corresponds to a **unique combination** of these factors.
-
-.. image:: teasers/dsprites_teaser.gif
- :align: center
- :width: 90%
-
Data Structure
--------------
@@ -70,50 +93,82 @@ When accessing an example using ``ds[i]``, you will receive a dictionary with th
- Description
* - ``image``
- ``PIL.Image.Image``
- - 64×64 binary image
+ - 64x64 (grayscale) or 64x64x3 (RGB) image
* - ``label``
- ``List[int]``
- Discrete latent indices: ``[color, shape, scale, orientation, posX, posY]``
* - ``label_values``
- ``List[float]``
- Continuous latent values corresponding to ``label``
- * - ``color`` … ``posY``
+ * - ``color`` ... ``posY``
- ``int``
- Individual discrete latent factors
- * - ``colorValue`` … ``posYValue``
+ * - ``colorValue`` ... ``posYValue``
- ``float``
- Individual continuous latent values
+ * - ``colorRGB``
+ - ``List[float]``
+ - Actual RGB color applied to the object (**color variant only**)
Usage Example
-------------
-**Basic Usage**
+**Basic Usage (original variant)**
.. code-block:: python
- from stable_datasets.images.dsprites import DSprites
-
- # First run will download + prepare cache, then return the split as a HF Dataset
- ds = DSprites(split="train")
+ from stable_datasets.images import DSprites
- # If you omit the split (split=None), you get a DatasetDict with all available splits
- ds_all = DSprites(split=None)
+ # Default variant is "original"
+ ds = DSprites(split="train", config_name="original")
sample = ds[0]
print(sample.keys())
- image = sample["image"]
- factors = sample["label"]
- factor_values = sample["label_values"]
+ image = sample["image"] # PIL.Image (64x64 grayscale)
+ factors = sample["label"] # [color, shape, scale, orientation, posX, posY]
# Optional: make it PyTorch-friendly
ds_torch = ds.with_format("torch")
+**Color variant**
+
+.. code-block:: python
+
+ from stable_datasets.images import DSprites
+
+ ds = DSprites(split="train", config_name="color")
+
+ sample = ds[0]
+ image = sample["image"] # PIL.Image (64x64x3 RGB)
+ color_rgb = sample["colorRGB"] # [R, G, B] in [0.5, 1.0]
+
+**Noise variant**
+
+.. code-block:: python
+
+ from stable_datasets.images import DSprites
+
+ ds = DSprites(split="train", config_name="noise")
+
+ sample = ds[0]
+ image = sample["image"] # PIL.Image (64x64x3, noisy background)
+
+**Scream variant**
+
+.. code-block:: python
+
+ from stable_datasets.images import DSprites
+
+ ds = DSprites(split="train", config_name="scream")
+
+ sample = ds[0]
+ image = sample["image"] # PIL.Image (64x64x3, Scream painting background)
Why No Train/Test Split?
------------------------
+------------------------
-The dSprites dataset does not define an official train/test split.
+The dSprites dataset does not define an official train/test split.
It is intended for **representation learning research**, where models are trained to capture underlying factors of variation rather than to generalize across semantic classes.
Because the dataset is a complete Cartesian product of all factor combinations, common evaluation protocols rely on:
@@ -122,22 +177,15 @@ Because the dataset is a complete Cartesian product of all factor combinations,
- Metric-based disentanglement scores
- Controlled interventions on latent variables
-Related Datasets
-----------------
-
-- **dSprites-Color**: Colored variant of dSprites
-- **dSprites-Noisy**: Noisy background variant
-- **dSprites-Scream**: Backgrounds replaced with natural images
-
References
----------
- Dataset repository: https://github.com/google-deepmind/dsprites-dataset
-- License: zlib/libpng License
-- Paper: Higgins et al., *β-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework*, ICLR 2017
+- Disentanglement library: https://github.com/google-research/disentanglement_lib/
+- License: zlib/libpng License (original), Apache License 2.0 (color/noise/scream)
-Citation
---------
+Citations
+---------
.. code-block:: bibtex
@@ -148,3 +196,15 @@ Citation
booktitle={International Conference on Learning Representations},
year={2017}
}
+
+.. code-block:: bibtex
+
+ @inproceedings{locatello2019challenging,
+ title={Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations},
+ author={Locatello, Francesco and Bauer, Stefan and Lucic, Mario and
+ Raetsch, Gunnar and Gelly, Sylvain and
+ Sch{\"o}lkopf, Bernhard and Bachem, Olivier},
+ booktitle={International Conference on Machine Learning},
+ pages={4114--4124},
+ year={2019}
+ }
diff --git a/docs/source/datasets/dsprites_color.rst b/docs/source/datasets/dsprites_color.rst
deleted file mode 100644
index 0f2750ed..00000000
--- a/docs/source/datasets/dsprites_color.rst
+++ /dev/null
@@ -1,163 +0,0 @@
-Color dSprites
-==============
-
-.. raw:: html
-
-
- 
- 
- 
-
-
-Overview
---------
-
-The Color dSprites dataset is a synthetic benchmark designed for **disentangled and unsupervised representation learning**. It is a variant of the original dSprites dataset, where the object in each image is rendered with a **random RGB color**, while the background remains black.
-
-Compared to the original grayscale dSprites, this variant introduces **random object color per sample**, enabling evaluation of **robustness to color variation** and analysis of how disentanglement methods behave under color transformations that are not part of the ground-truth generative factors.
-
-The dataset contains **all possible combinations** of six latent factors of variation inherited from dSprites, with each combination appearing exactly once.
-
-- **Total images**: 737,280
-- **Image resolution**: 64×64x3 (RGB)
-
-Latent Factors of Variation
---------------------------
-
-The dataset is generated from six independent latent factors, consistent with the original dSprites specification.
-In this color variant, the ``color`` factor remains fixed in the labels, while the **actual applied RGB color** is provided separately.
-
-.. list-table::
- :header-rows: 1
- :widths: 20 30 30
-
- * - Factor
- - Discrete Values
- - Continuous / Actual Values
- * - ``color``
- - {0}
- - Fixed label; actual RGB stored in ``colorRGB`` ∈ [0.5, 1.0]
- * - ``shape``
- - {0, 1, 2}
- - {1.0, 2.0, 3.0} (square, ellipse, heart)
- * - ``scale``
- - {0, ..., 5}
- - Linearly spaced in [0.5, 1.0]
- * - ``orientation``
- - {0, ..., 39}
- - Uniform in [0, 2π] radians
- * - ``posX``
- - {0, ..., 31}
- - Normalized position in [0, 1]
- * - ``posY``
- - {0, ..., 31}
- - Normalized position in [0, 1]
-
-Each image corresponds to a **unique combination** of these six latent factors.
-The random object color does **not** alter the factor indexing order.
-
-.. image:: teasers/dsprites_color_teaser.gif
- :align: center
- :width: 90%
-
-Data Structure
---------------
-
-When accessing an example using ``ds[i]``, you will receive a dictionary with the following keys:
-
-.. list-table::
- :header-rows: 1
- :widths: 25 20 55
-
- * - Key
- - Type
- - Description
- * - ``image``
- - ``PIL.Image.Image``
- - 64×64×3 RGB image (black background)
- * - ``label``
- - ``List[int]``
- - Discrete latent indices: ``[color, shape, scale, orientation, posX, posY]``
- * - ``label_values``
- - ``List[float]``
- - Continuous latent values corresponding to ``label``
- * - ``colorRGB``
- - ``List[float]``
- - Actual RGB color applied to the object
- * - ``color`` … ``posY``
- - ``int``
- - Individual discrete latent factors
- * - ``colorValue`` … ``posYValue``
- - ``float``
- - Individual continuous latent values
-
-**Note:**
-The ``color`` and ``colorValue`` fields remain fixed (0 / 1.0) to preserve compatibility with the original dSprites format.
-The actual object color is provided exclusively via ``colorRGB``.
-
-Usage Example
--------------
-
-**Basic Usage**
-
-.. code-block:: python
-
- from stable_datasets.images.dsprites_color import DSpritesColor
-
- # First run will download + prepare cache, then return the split as a HF Dataset
- ds = DSpritesColor(split="train")
-
- # If you omit the split (split=None), you get a DatasetDict with all available splits
- ds_all = DSpritesColor(split=None)
-
- sample = ds[0]
- print(sample.keys())
-
- image = sample["image"]
- factors = sample["label"]
- factor_values = sample["label_values"]
- color_rgb = sample["colorRGB"]
-
- # Optional: make it PyTorch-friendly
- ds_torch = ds.with_format("torch")
-
-Why No Train/Test Split?
------------------------
-
-The Color dSprites dataset does not define an official train/test split.
-It is intended for **representation learning research**, where models are trained to capture underlying factors of variation rather than to generalize across semantic classes.
-
-Because the dataset is a complete Cartesian product of all factor combinations, common evaluation protocols rely on:
-
-- Factor-wise generalization
-- Metric-based disentanglement scores
-- Robustness analysis under nuisance variations (e.g., color)
-- Controlled interventions on latent variables
-
-Related Datasets
-----------------
-
-- **dSprites**: Original grayscale version
-- **dSprites-Noisy**: Variant with background noise
-- **dSprites-Scream**: Variant with natural image backgrounds
-
-References
-----------
-
-- Dataset repository: https://github.com/google-research/disentanglement_lib/
-- License: Apache License 2.0
-- Paper: Locatello et al., *Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations*, ICML 2019
-
-Citation
---------
-
-.. code-block:: bibtex
-
- @inproceedings{locatello2019challenging,
- title={Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations},
- author={Locatello, Francesco and Bauer, Stefan and Lucic, Mario and
- Raetsch, Gunnar and Gelly, Sylvain and
- Sch{\"o}lkopf, Bernhard and Bachem, Olivier},
- booktitle={International Conference on Machine Learning},
- year={2019}
- }
diff --git a/docs/source/datasets/dsprites_noise.rst b/docs/source/datasets/dsprites_noise.rst
deleted file mode 100644
index 4655bd44..00000000
--- a/docs/source/datasets/dsprites_noise.rst
+++ /dev/null
@@ -1,159 +0,0 @@
-Noisy dSprites
-==============
-
-.. raw:: html
-
-
-
-
-
-
-
-Overview
---------
-
-The Noisy dSprites dataset is a synthetic benchmark designed for **disentangled and unsupervised representation learning**. It is a variant of the original dSprites dataset, where **random noise is added to the background** of each image, while the object itself remains unchanged.
-
-Compared to the original grayscale dSprites, this variant introduces **background noise as a nuisance factor**, enabling evaluation of **robustness to noisy observations** and analysis of how disentanglement methods perform when irrelevant background variation is present.
-
-The dataset contains **all possible combinations** of six latent factors of variation inherited from dSprites, with each combination appearing exactly once.
-
-- **Total images**: 737,280
-- **Image resolution**: 64×64x3 (binary object with noisy background)
-
-Latent Factors of Variation
---------------------------
-
-The dataset is generated from six independent latent factors, consistent with the original dSprites specification.
-In this noisy variant, **no additional latent factor is introduced**; background noise is applied independently of the ground-truth factors.
-
-.. list-table::
- :header-rows: 1
- :widths: 20 30 30
-
- * - Factor
- - Discrete Values
- - Continuous Values
- * - ``color``
- - {0}
- - 1.0 (fixed, white)
- * - ``shape``
- - {0, 1, 2}
- - {1.0, 2.0, 3.0} (square, ellipse, heart)
- * - ``scale``
- - {0, ..., 5}
- - Linearly spaced in [0.5, 1.0]
- * - ``orientation``
- - {0, ..., 39}
- - Uniform in [0, 2π] radians
- * - ``posX``
- - {0, ..., 31}
- - Normalized position in [0, 1]
- * - ``posY``
- - {0, ..., 31}
- - Normalized position in [0, 1]
-
-Each image corresponds to a **unique combination** of these six latent factors.
-The background noise does **not** affect the factor indexing or ordering.
-
-.. image:: teasers/dsprites_noisy_teaser.gif
- :align: center
- :width: 90%
-
-Data Structure
---------------
-
-When accessing an example using ``ds[i]``, you will receive a dictionary with the following keys:
-
-.. list-table::
- :header-rows: 1
- :widths: 25 20 55
-
- * - Key
- - Type
- - Description
- * - ``image``
- - ``PIL.Image.Image``
- - 64×64×3 image with noisy background
- * - ``label``
- - ``List[int]``
- - Discrete latent indices: ``[color, shape, scale, orientation, posX, posY]``
- * - ``label_values``
- - ``List[float]``
- - Continuous latent values corresponding to ``label``
- * - ``color`` … ``posY``
- - ``int``
- - Individual discrete latent factors
- * - ``colorValue`` … ``posYValue``
- - ``float``
- - Individual continuous latent values
-
-**Note:**
-In this Noisy variant, the object remains white and identical to the original dSprites.
-Only the **background pixels are replaced with random noise**, applied independently per image.
-
-Usage Example
--------------
-
-**Basic Usage**
-
-.. code-block:: python
-
- from stable_datasets.images.dsprites_noisy import DSpritesNoisy
-
- # First run will download + prepare cache, then return the split as a HF Dataset
- ds = DSpritesNoisy(split="train")
-
- # If you omit the split (split=None), you get a DatasetDict with all available splits
- ds_all = DSpritesNoisy(split=None)
-
- sample = ds[0]
- print(sample.keys())
-
- image = sample["image"]
- factors = sample["label"]
- factor_values = sample["label_values"]
-
- # Optional: make it PyTorch-friendly
- ds_torch = ds.with_format("torch")
-
-Why No Train/Test Split?
------------------------
-
-The Noisy dSprites dataset does not define an official train/test split.
-It is intended for **representation learning research**, where models are trained to capture underlying factors of variation rather than to generalize across semantic classes.
-
-Because the dataset is a complete Cartesian product of all factor combinations, common evaluation protocols rely on:
-
-- Factor-wise generalization
-- Metric-based disentanglement scores
-- Robustness to background noise
-- Controlled interventions on latent variables
-
-Related Datasets
-----------------
-
-- **dSprites**: Original grayscale version
-- **dSprites-Color**: Variant with random object color
-- **dSprites-Scream**: Variant with natural image backgrounds
-
-References
-----------
-
-- Dataset repository: https://github.com/google-research/disentanglement_lib/
-- License: Apache License 2.0
-- Paper: Locatello et al., *Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations*, ICML 2019
-
-Citation
---------
-
-.. code-block:: bibtex
-
- @inproceedings{locatello2019challenging,
- title={Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations},
- author={Locatello, Francesco and Bauer, Stefan and Lucic, Mario and
- Raetsch, Gunnar and Gelly, Sylvain and
- Sch{\"o}lkopf, Bernhard and Bachem, Olivier},
- booktitle={International Conference on Machine Learning},
- year={2019}
- }
diff --git a/docs/source/datasets/dsprites_scream.rst b/docs/source/datasets/dsprites_scream.rst
deleted file mode 100644
index 4107c72f..00000000
--- a/docs/source/datasets/dsprites_scream.rst
+++ /dev/null
@@ -1,168 +0,0 @@
-Scream dSprites
-===============
-
-.. raw:: html
-
-
-
-
-
-
-
-Overview
---------
-
-The Scream dSprites dataset is a synthetic benchmark designed for **disentangled and unsupervised representation learning**. It is a variant of the original dSprites dataset, where each object sprite is embedded onto a **random patch from a Scream painting image** used as background.
-
-In this Scream variant, the object is rendered by **inverting the pixel colors** of the sprite region on top of the background patch. This introduces a **structured and highly textured background**, enabling evaluation of robustness to complex background variation and testing how well models can disentangle latent factors under realistic nuisance conditions.
-
-The dataset contains **all possible combinations** of six latent factors of variation inherited from dSprites, with each combination appearing exactly once.
-
-- **Total images**: 737,280
-- **Image resolution**: 64×64×3 (textured background with inverted object)
-
-Latent Factors of Variation
---------------------------
-
-The dataset is generated from six independent latent factors, consistent with the original dSprites specification.
-In this variant, no additional latent factor is introduced; the background texture is applied independently of the ground-truth factors.
-
-.. list-table::
- :header-rows: 1
- :widths: 20 30 30
-
- * - Factor
- - Discrete Values
- - Description
- * - ``color``
- - {0}
- - Fixed label; object pixels inverted on Scream background
- * - ``shape``
- - {0, 1, 2}
- - Square, ellipse, heart
- * - ``scale``
- - {0, ..., 5}
- - Linearly spaced in [0.5, 1.0]
- * - ``orientation``
- - {0, ..., 39}
- - Uniform in [0, 2π] radians
- * - ``posX``
- - {0, ..., 31}
- - Normalized position in [0, 1]
- * - ``posY``
- - {0, ..., 31}
- - Normalized position in [0, 1]
-
-Each image corresponds to a **unique combination** of these six latent factors.
-The factor indexing and ordering are identical to the original dSprites dataset.
-
-.. image:: teasers/dsprites_scream_teaser.gif
- :align: center
- :width: 90%
-
-Scream Background Visualization
--------------------------------
-
-The background patches are sampled from a fixed Scream painting image.
-
-.. image:: imgs/scream.png
- :align: center
- :width: 70%
-
-Data Structure
---------------
-
-When accessing an example using ``ds[i]``, you will receive a dictionary with the following keys:
-
-.. list-table::
- :header-rows: 1
- :widths: 25 20 55
-
- * - Key
- - Type
- - Description
- * - ``image``
- - ``PIL.Image.Image``
- - 64×64×3 image with Scream background
- * - ``label``
- - ``List[int]``
- - Discrete latent indices: ``[color, shape, scale, orientation, posX, posY]``
- * - ``label_values``
- - ``List[float]``
- - Continuous latent values corresponding to ``label``
- * - ``color`` … ``posY``
- - ``int``
- - Individual discrete latent factors
- * - ``colorValue`` … ``posYValue``
- - ``float``
- - Individual continuous latent values
-
-**Note:**
-The ``color`` and ``colorValue`` fields remain fixed to preserve compatibility with the original dSprites format.
-The object is rendered by **pixel inversion** on top of the Scream background.
-
-Usage Example
--------------
-
-**Basic Usage**
-
-.. code-block:: python
-
- from stable_datasets.images.dsprites_scream import DSpritesScream
-
- # First run will download + prepare cache, then return the split as a HF Dataset
- ds = DSpritesScream(split="train")
-
- # If you omit the split (split=None), you get a DatasetDict with all available splits
- ds_all = DSpritesScream(split=None)
-
- sample = ds[0]
- print(sample.keys())
-
- image = sample["image"]
- factors = sample["label"]
- factor_values = sample["label_values"]
-
- # Optional: make it PyTorch-friendly
- ds_torch = ds.with_format("torch")
-
-Why No Train/Test Split?
------------------------
-
-The Scream dSprites dataset does not define an official train/test split.
-It is intended for **representation learning research**, where models are trained to capture underlying factors of variation rather than to generalize across semantic classes.
-
-Because the dataset is a complete Cartesian product of all factor combinations, common evaluation protocols rely on:
-
-- Factor-wise generalization
-- Metric-based disentanglement scores
-- Robustness to structured background textures
-- Controlled interventions on latent variables
-
-Related Datasets
-----------------
-
-- **dSprites**: Original grayscale version
-- **dSprites-Color**: Variant with random object color
-- **dSprites-Noisy**: Variant with random background noise
-
-References
-----------
-
-- Dataset repository: https://github.com/google-research/disentanglement_lib/
-- License: Apache License 2.0
-- Paper: Locatello et al., *Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations*, ICML 2019
-
-Citation
---------
-
-.. code-block:: bibtex
-
- @inproceedings{locatello2019challenging,
- title={Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations},
- author={Locatello, Francesco and Bauer, Stefan and Lucic, Mario and
- Raetsch, Gunnar and Gelly, Sylvain and
- Sch{\"o}lkopf, Bernhard and Bachem, Olivier},
- booktitle={International Conference on Machine Learning},
- year={2019}
- }
diff --git a/docs/source/datasets/eurosat.rst b/docs/source/datasets/eurosat.rst
new file mode 100644
index 00000000..521893fb
--- /dev/null
+++ b/docs/source/datasets/eurosat.rst
@@ -0,0 +1,123 @@
+EuroSAT
+=======
+
+.. raw:: html
+
+
+ 
+ 
+
+
+
+Overview
+--------
+EuroSAT is a land use and land cover classification benchmark built from Sentinel-2 satellite imagery. The RGB version contains 27,000 labeled 64×64 JPEG patches across 10 classes covering crops, forest, vegetation, industrial, residential, water, and transport categories.
+
+The original release does not ship an official split. This builder uses the widely adopted ``google-research/remote_sensing_representations`` split (16,200 train / 5,400 validation / 5,400 test), as mirrored by the ``timm/eurosat-rgb`` HuggingFace dataset in Parquet form.
+
+.. image:: teasers/eurosat_teaser.png
+ :align: center
+ :width: 90%
+
+
+- **Train**: 16,200 images
+- **Validation**: 5,400 images
+- **Test**: 5,400 images
+
+Classes (total 27,000 images)
+-----------------------------
+
+.. list-table::
+ :header-rows: 1
+ :widths: 30 20 50
+
+ * - Class
+ - # Images
+ - Description
+ * - ``AnnualCrop``
+ - 3,000
+ - Fields cultivated annually
+ * - ``Forest``
+ - 3,000
+ - Forest cover
+ * - ``HerbaceousVegetation``
+ - 3,000
+ - Grasslands and other herbaceous cover
+ * - ``Highway``
+ - 2,500
+ - Highways and major roads
+ * - ``Industrial``
+ - 2,500
+ - Industrial buildings / sites
+ * - ``Pasture``
+ - 2,000
+ - Pasture land
+ * - ``PermanentCrop``
+ - 2,500
+ - Permanent crops (e.g. orchards, vineyards)
+ * - ``Residential``
+ - 3,000
+ - Residential buildings
+ * - ``River``
+ - 2,500
+ - Rivers
+ * - ``SeaLake``
+ - 3,000
+ - Seas and lakes
+
+Data Structure
+--------------
+
+When accessing an example using ``ds[i]``, you will receive a dictionary with the following keys:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 20 20 60
+
+ * - Key
+ - Type
+ - Description
+ * - ``image``
+ - ``PIL.Image.Image``
+ - 64×64×3 RGB image
+ * - ``label``
+ - int
+ - Class label (0-9)
+
+Usage Example
+-------------
+
+**Basic Usage**
+
+.. code-block:: python
+
+ from stable_datasets.images.eurosat import EuroSAT
+
+ ds_train = EuroSAT(split="train")
+ ds_val = EuroSAT(split="validation")
+ ds_test = EuroSAT(split="test")
+
+ sample = ds_train[0]
+ print(sample.keys()) # {"image", "label"}
+
+ # Optional: make it PyTorch-friendly
+ ds_torch = ds_train.with_format("torch")
+
+References
+----------
+
+- Official repository: https://github.com/phelber/EuroSAT
+- License: MIT License
+
+Citation
+--------
+
+.. code-block:: bibtex
+
+ @article{helber2019eurosat,
+ title={Eurosat: A novel dataset and deep learning benchmark for land use and land cover classification},
+ author={Helber, Patrick and Bischke, Benjamin and Dengel, Andreas and Borth, Damian},
+ journal={IEEE Journal of Selected Topics in Applied Earth Observations and Remote Sensing},
+ year={2019},
+ publisher={IEEE}
+ }
diff --git a/docs/source/datasets/index.rst b/docs/source/datasets/index.rst
index 39de52d5..246e511e 100644
--- a/docs/source/datasets/index.rst
+++ b/docs/source/datasets/index.rst
@@ -51,6 +51,10 @@ Available Datasets
cifar10_c
cifar100_c
cars196
+ oxford_pet
+ plant_village
+ eurosat
+ stanford_dogs
dtd
fashion_mnist
k_mnist
diff --git a/docs/source/datasets/medmnist.rst b/docs/source/datasets/medmnist.rst
index 2f3a86be..eec54f69 100644
--- a/docs/source/datasets/medmnist.rst
+++ b/docs/source/datasets/medmnist.rst
@@ -14,7 +14,7 @@ Overview
MedMNIST is a large-scale MNIST-like collection of standardized biomedical images. In stable-datasets, MedMNIST is exposed via the `MedMNIST` class, but **the actual dataset depends on the selected variant** (passed as ``config_name``, e.g. ``dermamnist``, ``pathmnist``). Each variant provides train/validation/test splits.
-All 2D variants are pre-processed to **28×28** images and all 3D variants are pre-processed to **28×28×28** volumes, with corresponding labels.
+All 2D variants default to **28×28** images (with optional sizes **64**, **128**, **224**), and all 3D variants default to **28×28×28** volumes (with optional size **64**). Pass ``size=`` to select a larger resolution.
.. image:: teasers/medmnist_dermamnist_teaser.png
:align: center
@@ -217,6 +217,9 @@ Usage Example
# Optional: make it PyTorch-friendly
ds_train_torch = ds_train.with_format("torch")
+ # Load a larger resolution (MedMNIST+): 64, 128, or 224 for 2D; 64 for 3D
+ ds_train_224 = MedMNIST(split="train", config_name=variant, size=224)
+
**Basic Usage (2D multi-label variant: chestmnist)**
.. code-block:: python
@@ -248,6 +251,9 @@ Usage Example
image = sample["image"] # nested list, shape (28, 28, 28)
label = sample["label"] # int
+ # Load a larger resolution (MedMNIST+): 64 for 3D
+ ds_train_64 = MedMNIST(split="train", config_name=variant, size=64)
+
References
----------
diff --git a/docs/source/datasets/oxford_pet.rst b/docs/source/datasets/oxford_pet.rst
new file mode 100644
index 00000000..8888a4d7
--- /dev/null
+++ b/docs/source/datasets/oxford_pet.rst
@@ -0,0 +1,80 @@
+Oxford-IIIT Pet
+===============
+
+.. raw:: html
+
+
+ 
+ 
+ 
+
+
+Overview
+--------
+The Oxford-IIIT Pet dataset is a 37-category pet image dataset with roughly 200 images per class, covering 12 cat breeds and 25 dog breeds. Images exhibit large variations in scale, pose, and lighting. The dataset ships with an official train/test split used in the original paper: 3,680 training images and 3,669 test images, for a total of 7,349 images. Image resolutions vary across samples. No resizing is applied by default.
+
+.. image:: teasers/oxford_pet_teaser.png
+ :align: center
+ :width: 90%
+
+
+- **Train**: 3680 images
+- **Test**: 3669 images
+
+Data Structure
+--------------
+
+When accessing an example using ``ds[i]``, you will receive a dictionary with the following keys:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 20 20 60
+
+ * - Key
+ - Type
+ - Description
+ * - ``image``
+ - ``PIL.Image.Image``
+ - H×W×3 RGB image
+ * - ``label``
+ - int
+ - Breed label (0-36)
+
+Usage Example
+-------------
+
+**Basic Usage**
+
+.. code-block:: python
+
+ from stable_datasets.images.oxford_pet import OxfordPet
+
+ # First run will download + prepare cache, then return the split as a HF Dataset
+ ds = OxfordPet(split="train")
+
+ # If you omit the split (split=None), you get a DatasetDict with all available splits
+ ds_all = OxfordPet(split=None)
+
+ sample = ds[0]
+ print(sample.keys()) # {"image", "label"}
+
+ # Optional: make it PyTorch-friendly
+ ds_torch = ds.with_format("torch")
+
+References
+----------
+
+- Official website: https://www.robots.ox.ac.uk/~vgg/data/pets/
+- License: Creative Commons Attribution-ShareAlike 4.0 International License
+
+Citation
+--------
+
+.. code-block:: bibtex
+
+ @InProceedings{parkhi12a,
+ author = "Omkar M. Parkhi and Andrea Vedaldi and Andrew Zisserman and C. V. Jawahar",
+ title = "Cats and Dogs",
+ booktitle = "IEEE Conference on Computer Vision and Pattern Recognition",
+ year = "2012",
+ }
diff --git a/docs/source/datasets/plant_village.rst b/docs/source/datasets/plant_village.rst
new file mode 100644
index 00000000..11ca32d0
--- /dev/null
+++ b/docs/source/datasets/plant_village.rst
@@ -0,0 +1,127 @@
+PlantVillage
+============
+
+.. raw:: html
+
+
+ 
+ 
+
+
+
+Overview
+--------
+The PlantVillage dataset is an open-access repository of roughly 54,000 leaf images covering 14 crop species and 26 plant diseases (plus healthy controls), for 38 classes in total. Each image is labelled with a ``Crop___Disease`` class name (e.g. ``Apple___Apple_scab``, ``Tomato___healthy``). The dataset is distributed in three variants — original ``color`` photographs, ``grayscale`` conversions, and ``segmented`` (background-removed) images — and each variant ships with a published ~80/20 train/test split that respects per-leaf grouping to avoid leakage.
+
+.. image:: teasers/plant_village_teaser.png
+ :align: center
+ :width: 90%
+
+
+Variants
+--------
+
+Select a variant with the ``config_name`` argument.
+
+.. list-table::
+ :header-rows: 1
+ :widths: 18 12 18 18 34
+
+ * - Variant
+ - Classes
+ - Train
+ - Test
+ - Description
+ * - ``color``
+ - 38
+ - 43,596
+ - 10,709
+ - Original RGB leaf photographs (default)
+ * - ``grayscale``
+ - 38
+ - 43,203
+ - 11,102
+ - Grayscale conversions of the color images
+ * - ``segmented``
+ - 38
+ - 42,984
+ - 11,322
+ - Leaves segmented from the background
+
+Data Structure
+--------------
+
+When accessing an example using ``ds[i]``, you will receive a dictionary with the following keys:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 20 20 60
+
+ * - Key
+ - Type
+ - Description
+ * - ``image``
+ - ``PIL.Image.Image``
+ - H×W×3 RGB image
+ * - ``label``
+ - int
+ - Class label (0-37); use ``ds.features["label"].int2str(label)`` to recover ``Crop___Disease``
+ * - ``crop``
+ - str
+ - Crop name (e.g. ``Apple``, ``Tomato``, ``Pepper,_bell``)
+ * - ``disease``
+ - str
+ - Disease name or ``healthy`` (e.g. ``Apple_scab``, ``Late_blight``)
+
+Usage Example
+-------------
+
+**Basic Usage**
+
+.. code-block:: python
+
+ from stable_datasets.images.plant_village import PlantVillage
+
+ # Default variant is "color"
+ ds = PlantVillage(split="train")
+
+ # Select a different variant explicitly
+ ds_gray = PlantVillage(split="train", config_name="grayscale")
+ ds_seg = PlantVillage(split="test", config_name="segmented")
+
+ # All splits at once
+ ds_all = PlantVillage(split=None, config_name="color")
+
+ sample = ds[0]
+ print(sample.keys()) # {"image", "label", "crop", "disease"}
+ print(sample["crop"], sample["disease"])
+
+ # Optional: make it PyTorch-friendly
+ ds_torch = ds.with_format("torch")
+
+References
+----------
+
+- Official repository: https://github.com/spMohanty/PlantVillage-Dataset
+- License: CC0 1.0 (public domain)
+
+Citation
+--------
+
+.. code-block:: bibtex
+
+ @article{hughes2015open,
+ title={An open access repository of images on plant health to enable the development of mobile disease diagnostics},
+ author={Hughes, David P. and Salath{\'e}, Marcel},
+ journal={arXiv preprint arXiv:1511.08060},
+ year={2015}
+ }
+
+ @article{mohanty2016using,
+ title={Using deep learning for image-based plant disease detection},
+ author={Mohanty, Sharada P. and Hughes, David P. and Salath{\'e}, Marcel},
+ journal={Frontiers in Plant Science},
+ volume={7},
+ pages={1419},
+ year={2016}
+ }
diff --git a/docs/source/datasets/shapes3d.rst b/docs/source/datasets/shapes3d.rst
index 3abb6be7..2e28845a 100644
--- a/docs/source/datasets/shapes3d.rst
+++ b/docs/source/datasets/shapes3d.rst
@@ -138,7 +138,7 @@ protocols rely on:
References
----------
-- Dataset homepage: https://github.com/google-deepmind/3dshapes-dataset
+- Dataset homepage: https://github.com/google-deepmind/3d-shapes
- License: Apache License 2.0
- Paper: Kim & Mnih, *Disentangling by Factorising*, ICML 2018
diff --git a/docs/source/datasets/stanford_dogs.rst b/docs/source/datasets/stanford_dogs.rst
new file mode 100644
index 00000000..1664a6ad
--- /dev/null
+++ b/docs/source/datasets/stanford_dogs.rst
@@ -0,0 +1,83 @@
+Stanford Dogs
+=============
+
+.. raw:: html
+
+
+ 
+ 
+
+
+
+Overview
+--------
+The Stanford Dogs dataset contains 20,580 images of 120 dog breeds, drawn from ImageNet and curated for fine-grained image classification. Each breed corresponds to one ImageNet synset (e.g. ``n02085620-Chihuahua``) and has roughly 150-250 images. The dataset ships with an official train/test split published in the original paper: 12,000 training images (100 per class) and 8,580 test images. Image resolutions vary across samples. No resizing is applied by default.
+
+.. image:: teasers/stanford_dogs_teaser.png
+ :align: center
+ :width: 90%
+
+
+- **Train**: 12,000 images
+- **Test**: 8,580 images
+
+Data Structure
+--------------
+
+When accessing an example using ``ds[i]``, you will receive a dictionary with the following keys:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 20 20 60
+
+ * - Key
+ - Type
+ - Description
+ * - ``image``
+ - ``PIL.Image.Image``
+ - H×W×3 RGB image
+ * - ``label``
+ - int
+ - Breed label (0-119)
+
+Usage Example
+-------------
+
+**Basic Usage**
+
+.. code-block:: python
+
+ from stable_datasets.images.stanford_dogs import StanfordDogs
+
+ # First run will download + prepare cache, then return the split as a HF Dataset
+ ds = StanfordDogs(split="train")
+
+ # If you omit the split (split=None), you get a DatasetDict with all available splits
+ ds_all = StanfordDogs(split=None)
+
+ sample = ds[0]
+ print(sample.keys()) # {"image", "label"}
+ print(ds.features["label"].int2str(sample["label"])) # e.g. "chihuahua"
+
+ # Optional: make it PyTorch-friendly
+ ds_torch = ds.with_format("torch")
+
+References
+----------
+
+- Official website: http://vision.stanford.edu/aditya86/ImageNetDogs/
+- License: For non-commercial research and educational purposes only
+
+Citation
+--------
+
+.. code-block:: bibtex
+
+ @inproceedings{KhoslaYaoJayadevaprakashFeiFei_FGVC2011,
+ author = {Aditya Khosla and Nityananda Jayadevaprakash and Bangpeng Yao and Li Fei-Fei},
+ title = {Novel Dataset for Fine-Grained Image Categorization},
+ booktitle = {First Workshop on Fine-Grained Visual Categorization, IEEE Conference on Computer Vision and Pattern Recognition},
+ year = {2011},
+ month = {June},
+ address = {Colorado Springs, CO}
+ }
diff --git a/docs/source/datasets/teasers/dsprites_color_teaser.gif b/docs/source/datasets/teasers/dsprites_color_teaser.gif
deleted file mode 100644
index 188ad193..00000000
Binary files a/docs/source/datasets/teasers/dsprites_color_teaser.gif and /dev/null differ
diff --git a/docs/source/datasets/teasers/dsprites_noisy_teaser.gif b/docs/source/datasets/teasers/dsprites_noisy_teaser.gif
deleted file mode 100644
index bf6b2a8b..00000000
Binary files a/docs/source/datasets/teasers/dsprites_noisy_teaser.gif and /dev/null differ
diff --git a/docs/source/datasets/teasers/dsprites_scream_teaser.gif b/docs/source/datasets/teasers/dsprites_scream_teaser.gif
deleted file mode 100644
index b7987cdf..00000000
Binary files a/docs/source/datasets/teasers/dsprites_scream_teaser.gif and /dev/null differ
diff --git a/docs/source/datasets/teasers/dsprites_teaser.gif b/docs/source/datasets/teasers/dsprites_teaser.gif
index e778a376..525571f7 100644
Binary files a/docs/source/datasets/teasers/dsprites_teaser.gif and b/docs/source/datasets/teasers/dsprites_teaser.gif differ
diff --git a/docs/source/datasets/teasers/eurosat_teaser.png b/docs/source/datasets/teasers/eurosat_teaser.png
new file mode 100644
index 00000000..19982a3e
Binary files /dev/null and b/docs/source/datasets/teasers/eurosat_teaser.png differ
diff --git a/docs/source/datasets/teasers/oxford_pet_teaser.png b/docs/source/datasets/teasers/oxford_pet_teaser.png
new file mode 100644
index 00000000..05c7de3a
Binary files /dev/null and b/docs/source/datasets/teasers/oxford_pet_teaser.png differ
diff --git a/docs/source/datasets/teasers/plant_village_teaser.png b/docs/source/datasets/teasers/plant_village_teaser.png
new file mode 100644
index 00000000..f449ded5
Binary files /dev/null and b/docs/source/datasets/teasers/plant_village_teaser.png differ
diff --git a/docs/source/datasets/teasers/stanford_dogs_teaser.png b/docs/source/datasets/teasers/stanford_dogs_teaser.png
new file mode 100644
index 00000000..23da8976
Binary files /dev/null and b/docs/source/datasets/teasers/stanford_dogs_teaser.png differ
diff --git a/stable_datasets/cache.py b/stable_datasets/cache.py
index d8e8db07..a73b8bad 100644
--- a/stable_datasets/cache.py
+++ b/stable_datasets/cache.py
@@ -103,9 +103,15 @@ def _features_fingerprint(features: Features) -> str:
return hashlib.sha256(repr(features).encode()).hexdigest()[:16]
-def cache_fingerprint(cls_name: str, version: str, config_name: str, split: str) -> str:
- """Deterministic cache directory name for a dataset variant + split."""
- key = f"{cls_name}:{version}:{config_name}:{split}"
+def cache_fingerprint(cls_name: str, version: str, config_name: str, split: str, extra: str = "") -> str:
+ """Deterministic cache directory name for a dataset variant + split.
+
+ ``extra`` discriminates runtime config overrides; empty preserves the legacy hash.
+ """
+ if extra:
+ key = f"{cls_name}:{version}:{config_name}:{split}:{extra}"
+ else:
+ key = f"{cls_name}:{version}:{config_name}:{split}"
digest = hashlib.sha256(key.encode()).hexdigest()[:16]
return f"{cls_name.lower()}_{config_name}_{split}_{digest}"
diff --git a/stable_datasets/images/__init__.py b/stable_datasets/images/__init__.py
index 2d2bb20c..a82e9d21 100644
--- a/stable_datasets/images/__init__.py
+++ b/stable_datasets/images/__init__.py
@@ -14,11 +14,9 @@
from .country211 import Country211
from .cub200 import CUB200
from .dsprites import DSprites
-from .dsprites_color import DSpritesColor
-from .dsprites_noise import DSpritesNoise
-from .dsprites_scream import DSpritesScream
from .dtd import DTD
from .e_mnist import EMNIST
+from .eurosat import EuroSAT
from .face_pointing import FacePointing
from .fashion_mnist import FashionMNIST
@@ -35,12 +33,15 @@
from .linnaeus5 import Linnaeus5
from .med_mnist import MedMNIST
from .not_mnist import NotMNIST
+from .oxford_pet import OxfordPet
+from .plant_village import PlantVillage
# from .mnist import MNIST
# from .places365_small import Places365Small
from .rock_paper_scissor import RockPaperScissor
from .shapes3d import Shapes3D
from .small_norb import SmallNORB
+from .stanford_dogs import StanfordDogs
from .stl10 import STL10
from .svhn import SVHN
from .tiny_imagenet import TinyImagenet
@@ -62,11 +63,9 @@
"Country211",
"CUB200",
"DSprites",
- "DSpritesColor",
- "DSpritesNoise",
- "DSpritesScream",
"DTD",
"EMNIST",
+ "EuroSAT",
"FacePointing",
"FashionMNIST",
"FGVCAircraft",
@@ -81,9 +80,12 @@
"Linnaeus5",
"MedMNIST",
"NotMNIST",
+ "OxfordPet",
+ "PlantVillage",
"RockPaperScissor",
"Shapes3D",
"SmallNORB",
+ "StanfordDogs",
"STL10",
"SVHN",
"TinyImagenet",
diff --git a/stable_datasets/images/dsprites.py b/stable_datasets/images/dsprites.py
index 7dc0e832..28948ce7 100644
--- a/stable_datasets/images/dsprites.py
+++ b/stable_datasets/images/dsprites.py
@@ -1,97 +1,207 @@
+import os
+
import numpy as np
from PIL import Image as PILImage
-from stable_datasets.schema import DatasetInfo, Features, Sequence, Value, Version
+from stable_datasets.schema import (
+ BuilderConfig,
+ DatasetInfo,
+ Features,
+ Sequence,
+ Value,
+ Version,
+)
from stable_datasets.schema import Image as ImageFeature
from stable_datasets.utils import BaseDatasetBuilder
-class DSprites(BaseDatasetBuilder):
- """DSprites
- dSprites is a dataset of 2D shapes procedurally generated from 6 ground truth independent latent factors. These factors are color, shape, scale, rotation, x and y positions of a sprite."""
+DSPRITES_VERSION = Version("1.0.0")
+
+_THIS_DIR = os.path.dirname(os.path.abspath(__file__))
+_PROJECT_ROOT = os.path.abspath(os.path.join(_THIS_DIR, "..", ".."))
+_SCREAM_PATH = os.path.join(_PROJECT_ROOT, "docs", "source", "datasets", "imgs", "scream.png")
- VERSION = Version("1.0.0")
+_DSPRITES_URL = "https://github.com/google-deepmind/dsprites-dataset/raw/refs/heads/master/dsprites_ndarray_co1sh3sc6or40x32y32_64x64.npz"
- SOURCE = {
- "homepage": "https://github.com/deepmind/dsprites-dataset",
- "assets": {
- "train": "https://github.com/google-deepmind/dsprites-dataset/raw/refs/heads/master/dsprites_ndarray_co1sh3sc6or40x32y32_64x64.npz",
- },
- "citation": """@inproceedings{higgins2017beta,
+_CITATION_HIGGINS = """@inproceedings{higgins2017beta,
title={beta-vae: Learning basic visual concepts with a constrained variational framework},
author={Higgins, Irina and Matthey, Loic and Pal, Arka and Burgess, Christopher and Glorot, Xavier and Botvinick, Matthew and Mohamed, Shakir and Lerchner, Alexander},
booktitle={International conference on learning representations},
- year={2017}""",
- }
+ year={2017}"""
+
+_CITATION_LOCATELLO = """@inproceedings{locatello2019challenging,
+ title={Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations},
+ author={Locatello, Francesco and Bauer, Stefan and Lucic, Mario and Raetsch, Gunnar and Gelly, Sylvain and Sch{\\\"o}lkopf, Bernhard and Bachem, Olivier},
+ booktitle={International Conference on Machine Learning},
+ pages={4114--4124},
+ year={2019}
+ }"""
+
+
+class DSpritesConfig(BuilderConfig):
+ """BuilderConfig for DSprites variants.
+
+ Args:
+ has_color_rgb: Whether this variant produces a ``colorRGB`` field.
+ """
+
+ def __init__(self, *, has_color_rgb: bool = False, **kwargs):
+ super().__init__(version=DSPRITES_VERSION, **kwargs)
+ self.has_color_rgb = has_color_rgb
+
+
+class DSprites(BaseDatasetBuilder):
+ """dSprites dataset family.
+
+ dSprites is a dataset of 2D shapes procedurally generated from 6 ground
+ truth independent latent factors. These factors are color, shape, scale,
+ rotation, x and y positions of a sprite.
+
+ Four variants are available via ``config_name``:
+
+ - ``original``: 64x64 binary grayscale images (default).
+ - ``color``: Object rendered with a random RGB color on a black background.
+ - ``noise``: White object on a random-noise RGB background.
+ - ``scream``: Object rendered by inverting pixels on a Scream painting patch.
+ """
+
+ VERSION = DSPRITES_VERSION
+
+ BUILDER_CONFIGS = [
+ DSpritesConfig(
+ name="original",
+ description="Original grayscale dSprites (64x64 binary black-and-white)",
+ ),
+ DSpritesConfig(
+ name="color",
+ description="Color variant: random RGB object on black background (64x64x3)",
+ has_color_rgb=True,
+ ),
+ DSpritesConfig(
+ name="noise",
+ description="Noise variant: white object on random-noise background (64x64x3)",
+ ),
+ DSpritesConfig(
+ name="scream",
+ description="Scream variant: object inverted on Scream painting patch (64x64x3)",
+ ),
+ ]
+
+ def _source(self) -> dict:
+ variant = self.config.name
+ citation = _CITATION_HIGGINS if variant in ("original", "scream") else _CITATION_LOCATELLO
+
+ return {
+ "homepage": "https://github.com/deepmind/dsprites-dataset",
+ "assets": {"train": _DSPRITES_URL},
+ "citation": citation,
+ }
def _info(self):
+ source = self._source()
+
+ features_dict = {
+ "image": ImageFeature(),
+ "index": Value("int32"),
+ "label": Sequence(Value("int32")),
+ "label_values": Sequence(Value("float32")),
+ "color": Value("int32"),
+ "shape": Value("int32"),
+ "scale": Value("int32"),
+ "orientation": Value("int32"),
+ "posX": Value("int32"),
+ "posY": Value("int32"),
+ "colorValue": Value("float64"),
+ }
+
+ if getattr(self.config, "has_color_rgb", False):
+ features_dict["colorRGB"] = Sequence(Value("float32"))
+
+ features_dict.update(
+ {
+ "shapeValue": Value("float64"),
+ "scaleValue": Value("float64"),
+ "orientationValue": Value("float64"),
+ "posXValue": Value("float64"),
+ "posYValue": Value("float64"),
+ }
+ )
+
return DatasetInfo(
- description=""""dSprites dataset: procedurally generated 2D shapes dataset with known ground-truth factors, "
- "commonly used for disentangled representation learning. "
- "Factors: color (1), shape (3), scale (6), orientation (40), position X (32), position Y (32). "
- "Images are 64x64 binary black-and-white.""",
- features=Features(
- {
- "image": ImageFeature(), # (64, 64), grayscale
- "index": Value("int32"), # index of the image
- "label": Sequence(Value("int32")), # 6 factor indices (classes)
- "label_values": Sequence(Value("float32")), # 6 factor continuous values
- "color": Value("int32"), # color index (always 0)
- "shape": Value("int32"), # shape index (0-2)
- "scale": Value("int32"), # scale index (0-5)
- "orientation": Value("int32"), # orientation index (0-39)
- "posX": Value("int32"), # posX index (0-31)
- "posY": Value("int32"), # posY index (0-31)
- "colorValue": Value("float64"), # color value (always 1.0)
- "shapeValue": Value("float64"), # shape value (1.0, 2.0, 3.0)
- "scaleValue": Value("float64"), # scale value (0.5, 1)
- "orientationValue": Value("float64"), # orientation value (0, 2pi)
- "posXValue": Value("float64"), # posX value (0, 1)
- "posYValue": Value("float64"), # posY value (0, 1)
- }
- ),
+ description=f"dSprites dataset ({self.config.name} variant).",
+ features=Features(features_dict),
supervised_keys=("image", "label"),
- homepage=self.SOURCE["homepage"],
- citation=self.SOURCE["citation"],
+ homepage=source["homepage"],
+ citation=source["citation"],
)
def _generate_examples(self, data_path, split):
- # Load npz
data = np.load(data_path, allow_pickle=True)
- images = data["imgs"] # shape: (737280, 64, 64), uint8
- latents_classes = data["latents_classes"] # shape: (737280, 6), int64
- latents_values = data["latents_values"] # shape: (737280, 6), float64
+ images = data["imgs"] # (737280, 64, 64), uint8
+ latents_classes = data["latents_classes"] # (737280, 6), int64
+ latents_values = data["latents_values"] # (737280, 6), float64
+
+ variant = self.config.name
+
+ # Pre-load scream background once
+ scream = None
+ if variant == "scream":
+ scream_img = PILImage.open(_SCREAM_PATH).convert("RGB")
+ scream_img = scream_img.resize((350, 274))
+ scream = np.array(scream_img).astype(np.float32) / 255.0
- # Iterate over images
for idx in range(len(images)):
img = images[idx] # (64, 64), uint8
- img = img * 255
- # Convert to PIL image, keep grayscale mode
- img_pil = PILImage.fromarray(img, mode="L")
-
- factors_classes = latents_classes[
- idx
- ].tolist() # [color_idx, shape_idx, scale_idx, orientation_idx, posX_idx, posY_idx]
+ factors_classes = latents_classes[idx].tolist()
factors_values = latents_values[idx].tolist()
- yield (
- idx,
- {
- "image": img_pil,
- "index": idx,
- "label": factors_classes,
- "label_values": factors_values,
- "color": factors_classes[0], # always 0
- "shape": factors_classes[1],
- "scale": factors_classes[2],
- "orientation": factors_classes[3],
- "posX": factors_classes[4],
- "posY": factors_classes[5],
- "colorValue": factors_values[0], # always 0.0
- "shapeValue": factors_values[1],
- "scaleValue": factors_values[2],
- "orientationValue": factors_values[3],
- "posXValue": factors_values[4],
- "posYValue": factors_values[5],
- },
- )
+ color_rgb = None
+
+ if variant == "original":
+ img_pil = PILImage.fromarray(img * 255, mode="L")
+
+ elif variant == "color":
+ img_f = img.astype(np.float32)
+ color_rgb = np.random.uniform(0.5, 1.0, size=(3,))
+ img_rgb = (img_f[..., None] * color_rgb) * 255
+ img_pil = PILImage.fromarray(img_rgb.astype(np.uint8), mode="RGB")
+
+ elif variant == "noise":
+ img_f = img.astype(np.float32)
+ noise = np.random.uniform(0, 1, size=(64, 64, 3))
+ img_rgb = np.minimum(img_f[..., None] + noise, 1.0) * 255
+ img_pil = PILImage.fromarray(img_rgb.astype(np.uint8), mode="RGB")
+
+ elif variant == "scream":
+ img_f = img.astype(np.float32)
+ x_crop = np.random.randint(0, scream.shape[0] - 64)
+ y_crop = np.random.randint(0, scream.shape[1] - 64)
+ background_patch = scream[x_crop : x_crop + 64, y_crop : y_crop + 64]
+ mask = img_f == 1
+ output_img = np.copy(background_patch)
+ output_img[mask] = 1.0 - background_patch[mask]
+ img_pil = PILImage.fromarray((output_img * 255).astype(np.uint8), mode="RGB")
+
+ example = {
+ "image": img_pil,
+ "index": idx,
+ "label": factors_classes,
+ "label_values": factors_values,
+ "color": factors_classes[0],
+ "shape": factors_classes[1],
+ "scale": factors_classes[2],
+ "orientation": factors_classes[3],
+ "posX": factors_classes[4],
+ "posY": factors_classes[5],
+ "colorValue": factors_values[0],
+ "shapeValue": factors_values[1],
+ "scaleValue": factors_values[2],
+ "orientationValue": factors_values[3],
+ "posXValue": factors_values[4],
+ "posYValue": factors_values[5],
+ }
+
+ if variant == "color":
+ example["colorRGB"] = color_rgb.tolist()
+
+ yield idx, example
diff --git a/stable_datasets/images/dsprites_color.py b/stable_datasets/images/dsprites_color.py
deleted file mode 100644
index 00d18ebb..00000000
--- a/stable_datasets/images/dsprites_color.py
+++ /dev/null
@@ -1,102 +0,0 @@
-import numpy as np
-from PIL import Image as PILImage
-
-from stable_datasets.schema import DatasetInfo, Features, Sequence, Value, Version
-from stable_datasets.schema import Image as ImageFeature
-from stable_datasets.utils import BaseDatasetBuilder
-
-
-class DSpritesColor(BaseDatasetBuilder):
- """DSprites
- dSprites is a dataset of 2D shapes procedurally generated from 6 ground truth independent latent factors. These factors are color, shape, scale, rotation, x and y positions of a sprite."""
-
- VERSION = Version("1.0.0")
-
- SOURCE = {
- "homepage": "https://github.com/deepmind/dsprites-dataset",
- "assets": {
- "train": "https://github.com/google-deepmind/dsprites-dataset/raw/refs/heads/master/dsprites_ndarray_co1sh3sc6or40x32y32_64x64.npz",
- },
- "citation": """@inproceedings{locatello2019challenging,
- title={Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations},
- author={Locatello, Francesco and Bauer, Stefan and Lucic, Mario and Raetsch, Gunnar and Gelly, Sylvain and Sch{\"o}lkopf, Bernhard and Bachem, Olivier},
- booktitle={International Conference on Machine Learning},
- pages={4114--4124},
- year={2019}
- }""",
- }
-
- def _info(self):
- return DatasetInfo(
- description=""""dSprites dataset: procedurally generated 2D shapes dataset with known ground-truth factors, "
- "commonly used for disentangled representation learning. "
- "Factors: color (1), shape (3), scale (6), orientation (40), position X (32), position Y (32). "
- "Images are 64x64x3 where the object is RGB and background is black""",
- features=Features(
- {
- "image": ImageFeature(), # (64, 64), grayscale
- "index": Value("int32"), # index of the image
- "label": Sequence(Value("int32")), # 6 factor indices (classes)
- "label_values": Sequence(Value("float32")), # 6 factor continuous values
- "color": Value("int32"), # color index (always 0)
- "shape": Value("int32"), # shape index (0-2)
- "scale": Value("int32"), # scale index (0-5)
- "orientation": Value("int32"), # orientation index (0-39)
- "posX": Value("int32"), # posX index (0-31)
- "posY": Value("int32"), # posY index (0-31)
- "colorValue": Value("float64"), # color value (always 1.0)
- "colorRGB": Sequence(Value("float32")), # color RGB values (0.5, 1.0)
- "shapeValue": Value("float64"), # shape value (1.0, 2.0, 3.0)
- "scaleValue": Value("float64"), # scale value (0.5, 1)
- "orientationValue": Value("float64"), # orientation value (0, 2pi)
- "posXValue": Value("float64"), # posX value (0, 1)
- "posYValue": Value("float64"), # posY value (0, 1)
- }
- ),
- supervised_keys=("image", "label"),
- homepage=self.SOURCE["homepage"],
- citation=self.SOURCE["citation"],
- )
-
- def _generate_examples(self, data_path, split):
- # Load npz
- data = np.load(data_path, allow_pickle=True)
- images = data["imgs"] # shape: (737280, 64, 64), uint8
- latents_classes = data["latents_classes"] # shape: (737280, 6), int64
- latents_values = data["latents_values"] # shape: (737280, 6), float64
-
- # Iterate over images
- for idx in range(len(images)):
- img = images[idx] # (64, 64), uint8
- img = img.astype(np.float32) / 1.0
- color_rgb = np.random.uniform(0.5, 1.0, size=(3,))
- img_rgb = (img[..., None] * color_rgb) * 255
- img_pil = PILImage.fromarray(img_rgb.astype(np.uint8), mode="RGB")
-
- factors_classes = latents_classes[
- idx
- ].tolist() # [color_idx, shape_idx, scale_idx, orientation_idx, posX_idx, posY_idx]
- factors_values = latents_values[idx].tolist()
-
- yield (
- idx,
- {
- "image": img_pil,
- "index": idx,
- "label": factors_classes,
- "label_values": factors_values,
- "color": factors_classes[0], # always 0
- "shape": factors_classes[1],
- "scale": factors_classes[2],
- "orientation": factors_classes[3],
- "posX": factors_classes[4],
- "posY": factors_classes[5],
- "colorValue": factors_values[0], # always 1.0
- "colorRGB": color_rgb.tolist(), # [R, G, B], ∈ [0.5,1.0]
- "shapeValue": factors_values[1],
- "scaleValue": factors_values[2],
- "orientationValue": factors_values[3],
- "posXValue": factors_values[4],
- "posYValue": factors_values[5],
- },
- )
diff --git a/stable_datasets/images/dsprites_noise.py b/stable_datasets/images/dsprites_noise.py
deleted file mode 100644
index 97fdee04..00000000
--- a/stable_datasets/images/dsprites_noise.py
+++ /dev/null
@@ -1,100 +0,0 @@
-import numpy as np
-from PIL import Image as PILImage
-
-from stable_datasets.schema import DatasetInfo, Features, Sequence, Value, Version
-from stable_datasets.schema import Image as ImageFeature
-from stable_datasets.utils import BaseDatasetBuilder
-
-
-class DSpritesNoise(BaseDatasetBuilder):
- """DSprites
- dSprites is a dataset of 2D shapes procedurally generated from 6 ground truth independent latent factors. These factors are color, shape, scale, rotation, x and y positions of a sprite."""
-
- VERSION = Version("1.0.0")
-
- SOURCE = {
- "homepage": "https://github.com/deepmind/dsprites-dataset",
- "assets": {
- "train": "https://github.com/google-deepmind/dsprites-dataset/raw/refs/heads/master/dsprites_ndarray_co1sh3sc6or40x32y32_64x64.npz",
- },
- "citation": """@inproceedings{locatello2019challenging,
- title={Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations},
- author={Locatello, Francesco and Bauer, Stefan and Lucic, Mario and Raetsch, Gunnar and Gelly, Sylvain and Sch{\"o}lkopf, Bernhard and Bachem, Olivier},
- booktitle={International Conference on Machine Learning},
- pages={4114--4124},
- year={2019}
- }""",
- }
-
- def _info(self):
- return DatasetInfo(
- description=""""dSprites dataset: procedurally generated 2D shapes dataset with known ground-truth factors, "
- "commonly used for disentangled representation learning. "
- "Factors: color (1), shape (3), scale (6), orientation (40), position X (32), position Y (32). "
- "Images are 64x64x3. The object is white and the background is noise""",
- features=Features(
- {
- "image": ImageFeature(), # (64, 64), grayscale
- "index": Value("int32"), # index of the image
- "label": Sequence(Value("int32")), # 6 factor indices (classes)
- "label_values": Sequence(Value("float32")), # 6 factor continuous values
- "color": Value("int32"), # color index (always 0)
- "shape": Value("int32"), # shape index (0-2)
- "scale": Value("int32"), # scale index (0-5)
- "orientation": Value("int32"), # orientation index (0-39)
- "posX": Value("int32"), # posX index (0-31)
- "posY": Value("int32"), # posY index (0-31)
- "colorValue": Value("float64"), # color value (always 1.0)
- "shapeValue": Value("float64"), # shape value (1.0, 2.0, 3.0)
- "scaleValue": Value("float64"), # scale value (0.5, 1)
- "orientationValue": Value("float64"), # orientation value (0, 2pi)
- "posXValue": Value("float64"), # posX value (0, 1)
- "posYValue": Value("float64"), # posY value (0, 1)
- }
- ),
- supervised_keys=("image", "label"),
- homepage=self.SOURCE["homepage"],
- citation=self.SOURCE["citation"],
- )
-
- def _generate_examples(self, data_path, split):
- # Load npz
- data = np.load(data_path, allow_pickle=True)
- images = data["imgs"] # shape: (737280, 64, 64), uint8
- latents_classes = data["latents_classes"] # shape: (737280, 6), int64
- latents_values = data["latents_values"] # shape: (737280, 6), float64
-
- # Iterate over images
- for idx in range(len(images)):
- img = images[idx] # (64, 64), uint8
- img = img.astype(np.float32) / 1.0
- noise = np.random.uniform(0, 1, size=(64, 64, 3))
- img_rgb = np.minimum(img[..., None] + noise, 1.0) * 255
- img_pil = PILImage.fromarray(img_rgb.astype(np.uint8), mode="RGB")
-
- factors_classes = latents_classes[
- idx
- ].tolist() # [color_idx, shape_idx, scale_idx, orientation_idx, posX_idx, posY_idx]
- factors_values = latents_values[idx].tolist()
-
- yield (
- idx,
- {
- "image": img_pil,
- "index": idx,
- "label": factors_classes,
- "label_values": factors_values,
- "color": factors_classes[0], # always 0
- "shape": factors_classes[1],
- "scale": factors_classes[2],
- "orientation": factors_classes[3],
- "posX": factors_classes[4],
- "posY": factors_classes[5],
- "colorValue": factors_values[0], # always 0.0
- "shapeValue": factors_values[1],
- "scaleValue": factors_values[2],
- "orientationValue": factors_values[3],
- "posXValue": factors_values[4],
- "posYValue": factors_values[5],
- },
- )
diff --git a/stable_datasets/images/dsprites_scream.py b/stable_datasets/images/dsprites_scream.py
deleted file mode 100644
index 05b085e6..00000000
--- a/stable_datasets/images/dsprites_scream.py
+++ /dev/null
@@ -1,123 +0,0 @@
-import os
-
-import numpy as np
-from PIL import Image as PILImage
-
-from stable_datasets.schema import DatasetInfo, Features, Sequence, Value, Version
-from stable_datasets.schema import Image as ImageFeature
-from stable_datasets.utils import BaseDatasetBuilder
-
-
-_THIS_DIR = os.path.dirname(os.path.abspath(__file__))
-
-_PROJECT_ROOT = os.path.abspath(os.path.join(_THIS_DIR, "..", ".."))
-
-scream_path = os.path.join(_PROJECT_ROOT, "docs", "source", "datasets", "imgs", "scream.png")
-
-
-class DSpritesScream(BaseDatasetBuilder):
- """DSprites
- dSprites is a dataset of 2D shapes procedurally generated from 6 ground truth independent latent factors. These factors are color, shape, scale, rotation, x and y positions of a sprite."""
-
- VERSION = Version("1.0.0")
-
- SOURCE = {
- "homepage": "https://github.com/deepmind/dsprites-dataset",
- "assets": {
- "train": "https://github.com/google-deepmind/dsprites-dataset/raw/refs/heads/master/dsprites_ndarray_co1sh3sc6or40x32y32_64x64.npz",
- },
- "citation": """@inproceedings{higgins2017beta,
- title={beta-vae: Learning basic visual concepts with a constrained variational framework},
- author={Higgins, Irina and Matthey, Loic and Pal, Arka and Burgess, Christopher and Glorot, Xavier and Botvinick, Matthew and Mohamed, Shakir and Lerchner, Alexander},
- booktitle={International conference on learning representations},
- year={2017}""",
- }
-
- def _info(self):
- return DatasetInfo(
- description=""""dSprites dataset: procedurally generated 2D shapes dataset with known ground-truth factors, "
- "commonly used for disentangled representation learning. "
- "Factors: color (1), shape (3), scale (6), orientation (40), position X (32), position Y (32). "
- "Images are 64x64 binary black-and-white.""",
- features=Features(
- {
- "image": ImageFeature(), # (64, 64), grayscale
- "index": Value("int32"), # index of the image
- "label": Sequence(Value("int32")), # 6 factor indices (classes)
- "label_values": Sequence(Value("float32")), # 6 factor continuous values
- "color": Value("int32"), # color index (always 0)
- "shape": Value("int32"), # shape index (0-2)
- "scale": Value("int32"), # scale index (0-5)
- "orientation": Value("int32"), # orientation index (0-39)
- "posX": Value("int32"), # posX index (0-31)
- "posY": Value("int32"), # posY index (0-31)
- "colorValue": Value("float64"), # color value (always 1.0)
- "shapeValue": Value("float64"), # shape value (1.0, 2.0, 3.0)
- "scaleValue": Value("float64"), # scale value (0.5, 1)
- "orientationValue": Value("float64"), # orientation value (0, 2pi)
- "posXValue": Value("float64"), # posX value (0, 1)
- "posYValue": Value("float64"), # posY value (0, 1)
- }
- ),
- supervised_keys=("image", "label"),
- homepage=self.SOURCE["homepage"],
- citation=self.SOURCE["citation"],
- )
-
- def _generate_examples(self, data_path, split):
- # Load dSprites data
- data = np.load(data_path, allow_pickle=True)
- images = data["imgs"] # shape: (737280, 64, 64), uint8
- latents_classes = data["latents_classes"] # shape: (737280, 6), int64
- latents_values = data["latents_values"] # shape: (737280, 6), float64
-
- # Load Scream image once
- scream_img = PILImage.open(scream_path).convert("RGB")
- scream_img = scream_img.resize((350, 274))
- scream = np.array(scream_img).astype(np.float32) / 255.0 # (H, W, 3)
-
- # Iterate over images
- for idx in range(len(images)):
- img = images[idx].astype(np.float32) # (64, 64), float32
-
- # Random scream patch
- x_crop = np.random.randint(0, scream.shape[0] - 64)
- y_crop = np.random.randint(0, scream.shape[1] - 64)
- background_patch = scream[x_crop : x_crop + 64, y_crop : y_crop + 64] # (64, 64, 3)
-
- # Create mask
- mask = img == 1
- mask = mask[..., None] # (64, 64, 1)
-
- # Invert object region
- output_img = np.copy(background_patch)
- output_img[mask.squeeze()] = 1.0 - background_patch[mask.squeeze()]
-
- # Convert to RGB PIL
- img_pil = PILImage.fromarray((output_img * 255).astype(np.uint8), mode="RGB")
-
- # Labels
- factors_classes = latents_classes[idx].tolist()
- factors_values = latents_values[idx].tolist()
-
- yield (
- idx,
- {
- "image": img_pil,
- "index": idx,
- "label": factors_classes,
- "label_values": factors_values,
- "color": factors_classes[0],
- "shape": factors_classes[1],
- "scale": factors_classes[2],
- "orientation": factors_classes[3],
- "posX": factors_classes[4],
- "posY": factors_classes[5],
- "colorValue": factors_values[0],
- "shapeValue": factors_values[1],
- "scaleValue": factors_values[2],
- "orientationValue": factors_values[3],
- "posXValue": factors_values[4],
- "posYValue": factors_values[5],
- },
- )
diff --git a/stable_datasets/images/eurosat.py b/stable_datasets/images/eurosat.py
new file mode 100644
index 00000000..bc4afe65
--- /dev/null
+++ b/stable_datasets/images/eurosat.py
@@ -0,0 +1,92 @@
+import io
+
+import pyarrow.parquet as pq
+from PIL import Image as PILImage
+from tqdm import tqdm
+
+from stable_datasets.schema import (
+ ClassLabel,
+ DatasetInfo,
+ Features,
+ Version,
+)
+from stable_datasets.schema import Image as ImageFeature
+from stable_datasets.utils import BaseDatasetBuilder
+
+
+class EuroSAT(BaseDatasetBuilder):
+ """EuroSAT RGB Dataset
+
+ EuroSAT is a land use and land cover classification benchmark built from Sentinel-2 satellite
+ imagery. The RGB version consists of 27,000 labeled 64x64 JPEG patches across 10 classes.
+ The original release does not provide an official split; this builder uses the widely adopted
+ ``google-research/remote_sensing_representations`` split (16,200 train / 5,400 validation /
+ 5,400 test) as mirrored by the ``timm/eurosat-rgb`` HuggingFace dataset.
+ """
+
+ VERSION = Version("1.0.0")
+
+ SOURCE = {
+ "homepage": "https://github.com/phelber/EuroSAT",
+ "assets": {
+ "train": "https://huggingface.co/datasets/timm/eurosat-rgb/resolve/main/data/train-00000-of-00001.parquet",
+ "val": "https://huggingface.co/datasets/timm/eurosat-rgb/resolve/main/data/validation-00000-of-00001.parquet",
+ "test": "https://huggingface.co/datasets/timm/eurosat-rgb/resolve/main/data/test-00000-of-00001.parquet",
+ },
+ "citation": """@article{helber2019eurosat,
+ title={Eurosat: A novel dataset and deep learning benchmark for land use and land cover classification},
+ author={Helber, Patrick and Bischke, Benjamin and Dengel, Andreas and Borth, Damian},
+ journal={IEEE Journal of Selected Topics in Applied Earth Observations and Remote Sensing},
+ year={2019},
+ publisher={IEEE}
+ }""",
+ }
+
+ def _info(self):
+ return DatasetInfo(
+ description=(
+ "EuroSAT RGB: 27,000 Sentinel-2 64x64 RGB patches covering 10 land use and land cover "
+ "classes, with the google-research split of 16,200 train / 5,400 validation / 5,400 test."
+ ),
+ features=Features(
+ {
+ "image": ImageFeature(),
+ "label": ClassLabel(names=self._labels()),
+ }
+ ),
+ supervised_keys=("image", "label"),
+ homepage=self.SOURCE["homepage"],
+ citation=self.SOURCE["citation"],
+ )
+
+ def _generate_examples(self, data_path, split):
+ """Read the parquet file for this split and yield (idx, example) pairs.
+
+ Each parquet row has an ``image`` struct ``{"bytes": ..., "path": ...}`` (HF Image encoding)
+ and an integer ``label``.
+ """
+ table = pq.read_table(data_path, columns=["image", "label"])
+ images = table.column("image").to_pylist()
+ labels = table.column("label").to_pylist()
+
+ for idx, (img_entry, label) in enumerate(
+ tqdm(zip(images, labels), total=len(images), desc=f"Processing EuroSAT {split}")
+ ):
+ raw = img_entry["bytes"] if isinstance(img_entry, dict) else img_entry
+ image = PILImage.open(io.BytesIO(raw)).convert("RGB")
+ yield idx, {"image": image, "label": int(label)}
+
+ @staticmethod
+ def _labels():
+ return [
+ "AnnualCrop",
+ "Forest",
+ "HerbaceousVegetation",
+ "Highway",
+ "Industrial",
+ "Pasture",
+ "PermanentCrop",
+ "Residential",
+ "River",
+ "SeaLake",
+ ]
diff --git a/stable_datasets/images/med_mnist.py b/stable_datasets/images/med_mnist.py
index 5b42cea4..51809b4e 100644
--- a/stable_datasets/images/med_mnist.py
+++ b/stable_datasets/images/med_mnist.py
@@ -16,15 +16,41 @@
MEDMNIST_VERSION = Version("1.0.0")
+_VALID_SIZES_2D = (28, 64, 128, 224)
+_VALID_SIZES_3D = (28, 64)
-class MedMNISTConfig(BuilderConfig):
- """BuilderConfig with per-variant metadata used by MedMNIST._info()."""
- def __init__(self, *, num_classes: int, is_3d: bool = False, multi_label: bool = False, **kwargs):
+class MedMNISTConfig(BuilderConfig):
+ """BuilderConfig with per-variant metadata used by MedMNIST._info().
+
+ Args:
+ num_classes: Number of target classes.
+ is_3d: Whether the variant is a 3D volumetric dataset.
+ multi_label: Whether the task is multi-label classification.
+ size: Image resolution. 2D datasets support 28, 64, 128, 224;
+ 3D datasets support 28, 64. Defaults to 28 (MNIST-like).
+ """
+
+ def __init__(
+ self,
+ *,
+ num_classes: int,
+ is_3d: bool = False,
+ multi_label: bool = False,
+ size: int = 28,
+ **kwargs,
+ ):
super().__init__(version=MEDMNIST_VERSION, **kwargs)
+ valid_sizes = _VALID_SIZES_3D if is_3d else _VALID_SIZES_2D
+ if size not in valid_sizes:
+ raise ValueError(
+ f"size={size} is not valid for {'3D' if is_3d else '2D'} variant "
+ f"'{kwargs.get('name', '?')}'. Choose from {valid_sizes}."
+ )
self.num_classes = num_classes
self.is_3d = is_3d
self.multi_label = multi_label
+ self.size = size
class MedMNIST(BaseDatasetBuilder):
@@ -61,27 +87,39 @@ class MedMNIST(BaseDatasetBuilder):
]
def _source(self) -> dict:
- """Variant-aware source definition (computed from self.config at runtime)."""
+ """Variant- and size-aware source definition."""
variant = self.config.name
- url = f"https://zenodo.org/records/10519652/files/{variant}.npz?download=1"
- # Single NPZ contains all splits; we map each split name to the same URL.
+ size = getattr(self.config, "size", 28)
+ is_3d = getattr(self.config, "is_3d", False)
+
+ valid_sizes = _VALID_SIZES_3D if is_3d else _VALID_SIZES_2D
+ if size not in valid_sizes:
+ raise ValueError(
+ f"size={size} is not valid for {'3D' if is_3d else '2D'} variant "
+ f"'{variant}'. Choose from {valid_sizes}."
+ )
+
+ filename = f"{variant}.npz" if size == 28 else f"{variant}_{size}.npz"
+ url = f"https://zenodo.org/records/10519652/files/{filename}?download=1"
+
return {
"homepage": "https://medmnist.com/",
"assets": {"train": url, "test": url, "val": url},
"citation": """@article{medmnistv2,
- title={MedMNIST v2-A large-scale lightweight benchmark for 2D and 3D biomedical image classification},
- author={Yang, Jiancheng and Shi, Rui and Wei, Donglai and Liu, Zequan and Zhao, Lin and Ke, Bilian and Pfister, Hanspeter and Ni, Bingbing},
- journal={Scientific Data},
- volume={10},
- number={1},
- pages={41},
- year={2023},
- publisher={Nature Publishing Group UK London}
- }""",
+ title={MedMNIST v2-A large-scale lightweight benchmark for 2D and 3D biomedical image classification},
+ author={Yang, Jiancheng and Shi, Rui and Wei, Donglai and Liu, Zequan and Zhao, Lin and Ke, Bilian and Pfister, Hanspeter and Ni, Bingbing},
+ journal={Scientific Data},
+ volume={10},
+ number={1},
+ pages={41},
+ year={2023},
+ publisher={Nature Publishing Group UK London}
+ }""",
}
def _info(self):
variant = self.config.name
+ size = getattr(self.config, "size", 28)
source = self._source()
if getattr(self.config, "multi_label", False): # multi-label instead of multi-class
@@ -89,13 +127,16 @@ def _info(self):
else:
label_feature = ClassLabel(num_classes=self.config.num_classes)
+ if getattr(self.config, "is_3d", False):
+ image_feature = Array3D(shape=(size, size, size), dtype="uint8")
+ else:
+ image_feature = Image()
+
return DatasetInfo(
- description=f"MedMNIST variant: {variant} dataset.",
+ description=f"MedMNIST variant: {variant} (size={size}).",
features=Features(
{
- "image": (
- Array3D(shape=(28, 28, 28), dtype="uint8") if getattr(self.config, "is_3d", False) else Image()
- ),
+ "image": image_feature,
"label": label_feature,
}
),
diff --git a/stable_datasets/images/oxford_pet.py b/stable_datasets/images/oxford_pet.py
new file mode 100644
index 00000000..1590d198
--- /dev/null
+++ b/stable_datasets/images/oxford_pet.py
@@ -0,0 +1,154 @@
+import io
+import os
+import tarfile
+
+from PIL import Image as PILImage
+from tqdm import tqdm
+
+from stable_datasets.schema import ClassLabel, DatasetInfo, Features, Version
+from stable_datasets.schema import Image as ImageFeature
+from stable_datasets.splits import Split, SplitGenerator
+from stable_datasets.utils import BaseDatasetBuilder, bulk_download
+
+
+class OxfordPet(BaseDatasetBuilder):
+ """Oxford-IIIT Pet Dataset
+
+ The Oxford-IIIT Pet dataset is a 37-category pet image dataset with roughly 200 images per class,
+ covering 12 cat breeds and 25 dog breeds. Images exhibit large variations in scale, pose, and
+ lighting. Each image is annotated with a breed label, a head bounding box, and a pixel-level
+ trimap segmentation mask. The dataset ships with an official train/test split used in the original
+ paper: 3,680 training images and 3,669 test images, for a total of 7,349 images.
+ """
+
+ VERSION = Version("1.0.0")
+
+ SOURCE = {
+ "homepage": "https://www.robots.ox.ac.uk/~vgg/data/pets/",
+ "assets": {
+ "images": "https://thor.robots.ox.ac.uk/~vgg/data/pets/images.tar.gz",
+ "annotations": "https://thor.robots.ox.ac.uk/~vgg/data/pets/annotations.tar.gz",
+ },
+ "citation": """@InProceedings{parkhi12a,
+ author = "Omkar M. Parkhi and Andrea Vedaldi and Andrew Zisserman and C. V. Jawahar",
+ title = "Cats and Dogs",
+ booktitle = "IEEE Conference on Computer Vision and Pattern Recognition",
+ year = "2012",
+ }""",
+ }
+
+ def _info(self):
+ return DatasetInfo(
+ description=(
+ "The Oxford-IIIT Pet dataset contains 7,349 images covering 37 pet breeds "
+ "(12 cats, 25 dogs) with roughly 200 images per class. The official split "
+ "provides 3,680 training images and 3,669 test images."
+ ),
+ features=Features(
+ {
+ "image": ImageFeature(),
+ "label": ClassLabel(names=self._labels()),
+ }
+ ),
+ supervised_keys=("image", "label"),
+ homepage=self.SOURCE["homepage"],
+ citation=self.SOURCE["citation"],
+ )
+
+ def _split_generators(self):
+ """Override default splitting because both the image tar and annotation tar
+ are needed to build any split."""
+ source = self._source()
+
+ key_url_map = {
+ "images": source["assets"]["images"],
+ "annotations": source["assets"]["annotations"],
+ }
+
+ urls = list(key_url_map.values())
+ local_paths = bulk_download(urls, dest_folder=self._raw_download_dir)
+ path_map = dict(zip(key_url_map.keys(), local_paths))
+
+ return [
+ SplitGenerator(
+ name=Split.TRAIN,
+ gen_kwargs={"path_map": path_map, "split": "train"},
+ ),
+ SplitGenerator(
+ name=Split.TEST,
+ gen_kwargs={"path_map": path_map, "split": "test"},
+ ),
+ ]
+
+ def _generate_examples(self, path_map, split):
+ """Yield (key, example) pairs for the requested split by joining split list with image tar."""
+ images_path = path_map["images"]
+ annotations_path = path_map["annotations"]
+
+ split_file = "trainval.txt" if split == "train" else "test.txt"
+
+ name_to_label = {}
+ with tarfile.open(annotations_path, "r:gz") as tar:
+ member = tar.getmember(f"annotations/{split_file}")
+ with tar.extractfile(member) as f:
+ for line in f.read().decode("utf-8").splitlines():
+ line = line.strip()
+ if not line or line.startswith("#"):
+ continue
+ parts = line.split()
+ image_name, class_id = parts[0], int(parts[1])
+ name_to_label[image_name] = class_id - 1
+
+ with tarfile.open(images_path, "r:gz") as tar:
+ for entry in tqdm(tar, desc=f"Processing {split} set"):
+ if not entry.isfile() or not entry.name.endswith(".jpg"):
+ continue
+ stem = os.path.basename(entry.name).rsplit(".", 1)[0]
+ if stem not in name_to_label:
+ continue
+ with tar.extractfile(entry) as f:
+ image = PILImage.open(io.BytesIO(f.read())).convert("RGB")
+ yield stem, {"image": image, "label": name_to_label[stem]}
+
+ @staticmethod
+ def _labels():
+ """37 breeds in the canonical CLASS-ID order from annotations/list.txt (1-indexed there, 0-indexed here)."""
+ return [
+ "abyssinian",
+ "american_bulldog",
+ "american_pit_bull_terrier",
+ "basset_hound",
+ "beagle",
+ "bengal",
+ "birman",
+ "bombay",
+ "boxer",
+ "british_shorthair",
+ "chihuahua",
+ "egyptian_mau",
+ "english_cocker_spaniel",
+ "english_setter",
+ "german_shorthaired",
+ "great_pyrenees",
+ "havanese",
+ "japanese_chin",
+ "keeshond",
+ "leonberger",
+ "maine_coon",
+ "miniature_pinscher",
+ "newfoundland",
+ "persian",
+ "pomeranian",
+ "pug",
+ "ragdoll",
+ "russian_blue",
+ "saint_bernard",
+ "samoyed",
+ "scottish_terrier",
+ "shiba_inu",
+ "siamese",
+ "sphynx",
+ "staffordshire_bull_terrier",
+ "wheaten_terrier",
+ "yorkshire_terrier",
+ ]
diff --git a/stable_datasets/images/plant_village.py b/stable_datasets/images/plant_village.py
new file mode 100644
index 00000000..81676768
--- /dev/null
+++ b/stable_datasets/images/plant_village.py
@@ -0,0 +1,216 @@
+import io
+from zipfile import ZipFile
+
+from PIL import Image as PILImage
+from tqdm import tqdm
+
+from stable_datasets.schema import (
+ BuilderConfig,
+ ClassLabel,
+ DatasetInfo,
+ Features,
+ Value,
+ Version,
+)
+from stable_datasets.schema import Image as ImageFeature
+from stable_datasets.splits import Split, SplitGenerator
+from stable_datasets.utils import BaseDatasetBuilder, bulk_download
+
+
+PLANT_VILLAGE_VERSION = Version("1.0.0")
+
+_VARIANTS = ("color", "grayscale", "segmented")
+
+_HF_REPO = "https://huggingface.co/datasets/mohanty/PlantVillage/resolve/main"
+
+_CLASS_NAMES = [
+ "Apple___Apple_scab",
+ "Apple___Black_rot",
+ "Apple___Cedar_apple_rust",
+ "Apple___healthy",
+ "Blueberry___healthy",
+ "Cherry_(including_sour)___Powdery_mildew",
+ "Cherry_(including_sour)___healthy",
+ "Corn_(maize)___Cercospora_leaf_spot Gray_leaf_spot",
+ "Corn_(maize)___Common_rust_",
+ "Corn_(maize)___Northern_Leaf_Blight",
+ "Corn_(maize)___healthy",
+ "Grape___Black_rot",
+ "Grape___Esca_(Black_Measles)",
+ "Grape___Leaf_blight_(Isariopsis_Leaf_Spot)",
+ "Grape___healthy",
+ "Orange___Haunglongbing_(Citrus_greening)",
+ "Peach___Bacterial_spot",
+ "Peach___healthy",
+ "Pepper,_bell___Bacterial_spot",
+ "Pepper,_bell___healthy",
+ "Potato___Early_blight",
+ "Potato___Late_blight",
+ "Potato___healthy",
+ "Raspberry___healthy",
+ "Soybean___healthy",
+ "Squash___Powdery_mildew",
+ "Strawberry___Leaf_scorch",
+ "Strawberry___healthy",
+ "Tomato___Bacterial_spot",
+ "Tomato___Early_blight",
+ "Tomato___Late_blight",
+ "Tomato___Leaf_Mold",
+ "Tomato___Septoria_leaf_spot",
+ "Tomato___Spider_mites Two-spotted_spider_mite",
+ "Tomato___Target_Spot",
+ "Tomato___Tomato_Yellow_Leaf_Curl_Virus",
+ "Tomato___Tomato_mosaic_virus",
+ "Tomato___healthy",
+]
+
+_CLASS_INDEX = {name: i for i, name in enumerate(_CLASS_NAMES)}
+
+
+class PlantVillageConfig(BuilderConfig):
+ """One BuilderConfig per image variant (color / grayscale / segmented)."""
+
+ def __init__(self, *, variant: str, **kwargs):
+ super().__init__(version=PLANT_VILLAGE_VERSION, **kwargs)
+ if variant not in _VARIANTS:
+ raise ValueError(f"variant={variant!r} is not valid; choose from {_VARIANTS}.")
+ self.variant = variant
+
+
+class PlantVillage(BaseDatasetBuilder):
+ """PlantVillage Dataset
+
+ The PlantVillage dataset is an open-access repository of leaf images covering 14 crop species
+ and 26 plant diseases (plus healthy controls), for 38 classes in total and ~54,000 images.
+ Three image variants are provided:
+
+ - ``color``: original RGB photographs of the leaves
+ - ``grayscale``: the same images converted to grayscale
+ - ``segmented``: leaves segmented from the background
+
+ The train/test split used here is the one published alongside the HuggingFace mirror
+ (``mohanty/PlantVillage``), with ~80/20 partitioning per class.
+ """
+
+ VERSION = PLANT_VILLAGE_VERSION
+
+ BUILDER_CONFIGS = [
+ PlantVillageConfig(
+ name="color",
+ description="Original RGB leaf photographs.",
+ variant="color",
+ ),
+ PlantVillageConfig(
+ name="grayscale",
+ description="Grayscale leaf photographs.",
+ variant="grayscale",
+ ),
+ PlantVillageConfig(
+ name="segmented",
+ description="Leaves segmented from the background.",
+ variant="segmented",
+ ),
+ ]
+ DEFAULT_CONFIG_NAME = "color"
+
+ def _source(self):
+ variant = self.config.variant
+ return {
+ "homepage": "https://github.com/spMohanty/PlantVillage-Dataset",
+ "assets": {
+ "data": f"{_HF_REPO}/data.zip",
+ "train_split": f"{_HF_REPO}/splits/{variant}_train.txt",
+ "test_split": f"{_HF_REPO}/splits/{variant}_test.txt",
+ },
+ "citation": """@article{hughes2015open,
+ title={An open access repository of images on plant health to enable the development of mobile disease diagnostics},
+ author={Hughes, David P. and Salath{\\'e}, Marcel},
+ journal={arXiv preprint arXiv:1511.08060},
+ year={2015}
+ }
+ @article{mohanty2016using,
+ title={Using deep learning for image-based plant disease detection},
+ author={Mohanty, Sharada P. and Hughes, David P. and Salath{\\'e}, Marcel},
+ journal={Frontiers in Plant Science},
+ volume={7},
+ pages={1419},
+ year={2016}
+ }""",
+ }
+
+ def _info(self):
+ source = self._source()
+ return DatasetInfo(
+ description=(
+ "PlantVillage plant-disease image dataset covering 14 crops and 26 diseases "
+ f"(38 classes total) in three image variants: {', '.join(_VARIANTS)}."
+ ),
+ features=Features(
+ {
+ "image": ImageFeature(),
+ "label": ClassLabel(names=_CLASS_NAMES),
+ "crop": Value("string"),
+ "disease": Value("string"),
+ }
+ ),
+ supervised_keys=("image", "label"),
+ homepage=source["homepage"],
+ citation=source["citation"],
+ )
+
+ def _split_generators(self):
+ source = self._source()
+ key_url_map = {
+ "data": source["assets"]["data"],
+ "train_split": source["assets"]["train_split"],
+ "test_split": source["assets"]["test_split"],
+ }
+ urls = list(key_url_map.values())
+ local_paths = bulk_download(urls, dest_folder=self._raw_download_dir)
+ path_map = dict(zip(key_url_map.keys(), local_paths))
+
+ return [
+ SplitGenerator(
+ name=Split.TRAIN,
+ gen_kwargs={"path_map": path_map, "split": "train"},
+ ),
+ SplitGenerator(
+ name=Split.TEST,
+ gen_kwargs={"path_map": path_map, "split": "test"},
+ ),
+ ]
+
+ def _generate_examples(self, path_map, split):
+ """For each entry in the split file, read its image from data.zip and derive crop/disease from the path."""
+ data_path = path_map["data"]
+ split_path = path_map["train_split"] if split == "train" else path_map["test_split"]
+
+ with open(split_path) as f:
+ rel_paths = [line.strip() for line in f if line.strip()]
+
+ with ZipFile(data_path, "r") as archive:
+ for rel_path in tqdm(rel_paths, desc=f"Processing {self.config.variant} {split} set"):
+ parts = rel_path.split("/")
+ # Expect: raw///
+ if len(parts) < 4:
+ continue
+ class_name = parts[2]
+ sub = class_name.split("___", 1)
+ crop = sub[0]
+ disease = sub[1] if len(sub) > 1 else "unknown"
+
+ try:
+ raw_bytes = archive.read(rel_path)
+ except KeyError:
+ continue
+ image = PILImage.open(io.BytesIO(raw_bytes)).convert("RGB")
+
+ yield (
+ rel_path,
+ {
+ "image": image,
+ "label": _CLASS_INDEX[class_name],
+ "crop": crop,
+ "disease": disease,
+ },
+ )
diff --git a/stable_datasets/images/shapes3d.py b/stable_datasets/images/shapes3d.py
index 9b195d89..fc0b7ae9 100644
--- a/stable_datasets/images/shapes3d.py
+++ b/stable_datasets/images/shapes3d.py
@@ -12,7 +12,7 @@ class Shapes3D(BaseDatasetBuilder):
VERSION = Version("1.0.0")
SOURCE = {
- "homepage": "https://github.com/google-deepmind/3dshapes-dataset/",
+ "homepage": "https://github.com/google-deepmind/3d-shapes",
"assets": {
"train": "https://huggingface.co/datasets/randall-lab/shapes3d/resolve/main/shapes3d.npz",
},
diff --git a/stable_datasets/images/stanford_dogs.py b/stable_datasets/images/stanford_dogs.py
new file mode 100644
index 00000000..de96b2d9
--- /dev/null
+++ b/stable_datasets/images/stanford_dogs.py
@@ -0,0 +1,231 @@
+import io
+import tarfile
+
+import scipy.io
+from PIL import Image as PILImage
+from tqdm import tqdm
+
+from stable_datasets.schema import ClassLabel, DatasetInfo, Features, Version
+from stable_datasets.schema import Image as ImageFeature
+from stable_datasets.splits import Split, SplitGenerator
+from stable_datasets.utils import BaseDatasetBuilder, bulk_download
+
+
+class StanfordDogs(BaseDatasetBuilder):
+ """Stanford Dogs Dataset
+
+ The Stanford Dogs dataset contains 20,580 images of 120 dog breeds, built on top of ImageNet.
+ Each breed has roughly 150-250 images and an associated synset id (e.g. ``n02085620-Chihuahua``).
+ The dataset ships with an official train/test split (12,000 train, 8,580 test) defined in
+ ``train_list.mat`` and ``test_list.mat``.
+ """
+
+ VERSION = Version("1.0.0")
+
+ SOURCE = {
+ "homepage": "http://vision.stanford.edu/aditya86/ImageNetDogs/",
+ "assets": {
+ "images": "http://vision.stanford.edu/aditya86/ImageNetDogs/images.tar",
+ "lists": "http://vision.stanford.edu/aditya86/ImageNetDogs/lists.tar",
+ },
+ "citation": """@inproceedings{KhoslaYaoJayadevaprakashFeiFei_FGVC2011,
+ author = {Aditya Khosla and Nityananda Jayadevaprakash and Bangpeng Yao and Li Fei-Fei},
+ title = {Novel Dataset for Fine-Grained Image Categorization},
+ booktitle = {First Workshop on Fine-Grained Visual Categorization, IEEE Conference on Computer Vision and Pattern Recognition},
+ year = {2011},
+ month = {June},
+ address = {Colorado Springs, CO}
+ }""",
+ }
+
+ def _info(self):
+ return DatasetInfo(
+ description=(
+ "Stanford Dogs: 20,580 images of 120 dog breeds drawn from ImageNet, with an "
+ "official 12,000 / 8,580 train/test split."
+ ),
+ features=Features(
+ {
+ "image": ImageFeature(),
+ "label": ClassLabel(names=self._labels()),
+ }
+ ),
+ supervised_keys=("image", "label"),
+ homepage=self.SOURCE["homepage"],
+ citation=self.SOURCE["citation"],
+ )
+
+ def _split_generators(self):
+ """Download both images.tar and lists.tar; each split joins them via the .mat split files."""
+ source = self._source()
+ key_url_map = {
+ "images": source["assets"]["images"],
+ "lists": source["assets"]["lists"],
+ }
+ urls = list(key_url_map.values())
+ local_paths = bulk_download(urls, dest_folder=self._raw_download_dir)
+ path_map = dict(zip(key_url_map.keys(), local_paths))
+
+ return [
+ SplitGenerator(
+ name=Split.TRAIN,
+ gen_kwargs={"path_map": path_map, "split": "train"},
+ ),
+ SplitGenerator(
+ name=Split.TEST,
+ gen_kwargs={"path_map": path_map, "split": "test"},
+ ),
+ ]
+
+ def _generate_examples(self, path_map, split):
+ """Read the .mat split file from lists.tar, then iterate images.tar once and yield the matched entries."""
+ images_path = path_map["images"]
+ lists_path = path_map["lists"]
+
+ split_mat_name = "train_list.mat" if split == "train" else "test_list.mat"
+ with tarfile.open(lists_path, "r") as lists_tar:
+ with lists_tar.extractfile(split_mat_name) as f:
+ mat = scipy.io.loadmat(io.BytesIO(f.read()))
+
+ # mat 'file_list' entries look like 'n02085620-Chihuahua/n02085620_5927.jpg' (relative to Images/)
+ # Labels in the mat are 1-indexed; convert to 0-indexed for the ClassLabel feature.
+ rel_files = [str(x[0][0]) for x in mat["file_list"]]
+ raw_labels = mat["labels"].ravel().tolist()
+ archive_path_to_label = {f"Images/{rel}": int(lab) - 1 for rel, lab in zip(rel_files, raw_labels)}
+
+ with tarfile.open(images_path, "r") as img_tar:
+ for entry in tqdm(img_tar, desc=f"Processing {split} set"):
+ if not entry.isfile() or not entry.name.endswith(".jpg"):
+ continue
+ label = archive_path_to_label.get(entry.name)
+ if label is None:
+ continue
+ with img_tar.extractfile(entry) as f:
+ image = PILImage.open(io.BytesIO(f.read())).convert("RGB")
+ yield entry.name, {"image": image, "label": label}
+
+ @staticmethod
+ def _labels():
+ """120 breed names in canonical CLASS-ID order from file_list.mat (1-indexed there, 0-indexed here)."""
+ return [
+ "chihuahua",
+ "japanese_spaniel",
+ "maltese_dog",
+ "pekinese",
+ "shih-tzu",
+ "blenheim_spaniel",
+ "papillon",
+ "toy_terrier",
+ "rhodesian_ridgeback",
+ "afghan_hound",
+ "basset",
+ "beagle",
+ "bloodhound",
+ "bluetick",
+ "black-and-tan_coonhound",
+ "walker_hound",
+ "english_foxhound",
+ "redbone",
+ "borzoi",
+ "irish_wolfhound",
+ "italian_greyhound",
+ "whippet",
+ "ibizan_hound",
+ "norwegian_elkhound",
+ "otterhound",
+ "saluki",
+ "scottish_deerhound",
+ "weimaraner",
+ "staffordshire_bullterrier",
+ "american_staffordshire_terrier",
+ "bedlington_terrier",
+ "border_terrier",
+ "kerry_blue_terrier",
+ "irish_terrier",
+ "norfolk_terrier",
+ "norwich_terrier",
+ "yorkshire_terrier",
+ "wire-haired_fox_terrier",
+ "lakeland_terrier",
+ "sealyham_terrier",
+ "airedale",
+ "cairn",
+ "australian_terrier",
+ "dandie_dinmont",
+ "boston_bull",
+ "miniature_schnauzer",
+ "giant_schnauzer",
+ "standard_schnauzer",
+ "scotch_terrier",
+ "tibetan_terrier",
+ "silky_terrier",
+ "soft-coated_wheaten_terrier",
+ "west_highland_white_terrier",
+ "lhasa",
+ "flat-coated_retriever",
+ "curly-coated_retriever",
+ "golden_retriever",
+ "labrador_retriever",
+ "chesapeake_bay_retriever",
+ "german_short-haired_pointer",
+ "vizsla",
+ "english_setter",
+ "irish_setter",
+ "gordon_setter",
+ "brittany_spaniel",
+ "clumber",
+ "english_springer",
+ "welsh_springer_spaniel",
+ "cocker_spaniel",
+ "sussex_spaniel",
+ "irish_water_spaniel",
+ "kuvasz",
+ "schipperke",
+ "groenendael",
+ "malinois",
+ "briard",
+ "kelpie",
+ "komondor",
+ "old_english_sheepdog",
+ "shetland_sheepdog",
+ "collie",
+ "border_collie",
+ "bouvier_des_flandres",
+ "rottweiler",
+ "german_shepherd",
+ "doberman",
+ "miniature_pinscher",
+ "greater_swiss_mountain_dog",
+ "bernese_mountain_dog",
+ "appenzeller",
+ "entlebucher",
+ "boxer",
+ "bull_mastiff",
+ "tibetan_mastiff",
+ "french_bulldog",
+ "great_dane",
+ "saint_bernard",
+ "eskimo_dog",
+ "malamute",
+ "siberian_husky",
+ "affenpinscher",
+ "basenji",
+ "pug",
+ "leonberg",
+ "newfoundland",
+ "great_pyrenees",
+ "samoyed",
+ "pomeranian",
+ "chow",
+ "keeshond",
+ "brabancon_griffon",
+ "pembroke",
+ "cardigan",
+ "toy_poodle",
+ "miniature_poodle",
+ "standard_poodle",
+ "mexican_hairless",
+ "dingo",
+ "dhole",
+ "african_hunting_dog",
+ ]
diff --git a/stable_datasets/tests/images/test_dsprites.py b/stable_datasets/tests/images/test_dsprites.py
index 02c807d4..db698fd4 100644
--- a/stable_datasets/tests/images/test_dsprites.py
+++ b/stable_datasets/tests/images/test_dsprites.py
@@ -4,82 +4,85 @@
from stable_datasets.images import DSprites
-def test_dsprites_dataset():
- # Load training split
- dsprites_train = DSprites(split="train")
+def test_dsprites_variants():
+ """Download/integration test for all DSprites variants."""
- # Test 1: Check number of training samples
- expected_num_train_samples = 737280
- assert len(dsprites_train) == expected_num_train_samples, (
- f"Expected {expected_num_train_samples} training samples, got {len(dsprites_train)}."
- )
+ variants = ["original", "color", "noise", "scream"]
- # Test 2: Check sample keys and label range
- sample = dsprites_train[0]
- expected_keys = {
- "image",
- "index",
- "label",
- "label_values",
- "color",
- "shape",
- "scale",
- "orientation",
- "posX",
- "posY",
- "colorValue",
- "shapeValue",
- "scaleValue",
- "orientationValue",
- "posXValue",
- "posYValue",
- }
- assert sorted(set(sample.keys())) == sorted(set(expected_keys)), (
- f"Expected keys {expected_keys}, got {set(sample.keys())}."
- )
+ for variant in variants:
+ ds = DSprites(split="train", config_name=variant)
- # Test 3: Validate image type
- image = sample["image"]
- assert isinstance(image, Image.Image), f"Image should be a PIL.Image.Image, got {type(image)}."
- image_np = np.array(image)
- assert image_np.ndim == 2, f"DSprites images should be HxW, got shape {image_np.shape}."
- assert image_np.dtype == np.uint8, f"Image dtype should be uint8, got {image_np.dtype}."
- assert image_np.shape == (64, 64), f"Image should have shape (64, 64), got {image_np.shape}"
+ # Check dataset size
+ expected_num_samples = 737280
+ assert len(ds) == expected_num_samples, f"{variant}: Expected {expected_num_samples} samples, got {len(ds)}."
- # Test 4: Validate label type and range
- label = sample["label"]
- label_values = sample["label_values"]
- assert isinstance(label, list), f"Label should be int, got {type(list)}."
- assert isinstance(label_values, list), f"Label values should be list, got {type(label_values)}."
- assert len(label) == 6, f"Label should have 6 elements, got {len(label)}."
- assert len(label_values) == 6, f"Label values should have 6 elements, got {len(label_values)}."
+ sample = ds[0]
- color = sample["color"]
- shape = sample["shape"]
- scale = sample["scale"]
- orientation = sample["orientation"]
- posX = sample["posX"]
- posY = sample["posY"]
- colorValue = sample["colorValue"]
- shapeValue = sample["shapeValue"]
- scaleValue = sample["scaleValue"]
- orientationValue = sample["orientationValue"]
- posXValue = sample["posXValue"]
- posYValue = sample["posYValue"]
+ # Check keys
+ expected_keys = {
+ "image",
+ "index",
+ "label",
+ "label_values",
+ "color",
+ "shape",
+ "scale",
+ "orientation",
+ "posX",
+ "posY",
+ "colorValue",
+ "shapeValue",
+ "scaleValue",
+ "orientationValue",
+ "posXValue",
+ "posYValue",
+ }
+ if variant == "color":
+ expected_keys.add("colorRGB")
- assert 0 <= color < 1, f"Color should be in range [0, 0], got {color}."
- assert 0 <= shape < 3, f"Shape should be in range [0, 2], got {shape}."
- assert 0 <= scale < 6, f"Scale should be in range [0, 5], got {scale}."
- assert 0 <= orientation < 40, f"Orientation should be in range [0, 39], got {orientation}."
- assert 0 <= posX < 32, f"PosX should be in range [0, 31], got {posX}."
- assert 0 <= posY < 32, f"PosY should be in range [0, 31], got {posY}."
- assert colorValue == 1.0, f"Color value should be 1.0, got {colorValue}."
- assert shapeValue in [1.0, 2.0, 3.0], f"Shape value should be in [1.0, 2.0, 3.0], got {shapeValue}."
- assert 0.5 <= scaleValue <= 1, f"Scale value should be in range [0.5, 1], got {scaleValue}."
- assert 0 <= orientationValue <= 2 * np.pi, (
- f"Orientation value should be in range [0, 2pi], got {orientationValue}."
- )
- assert 0 <= posXValue <= 1, f"PosX value should be in range [0, 1], got {posXValue}."
- assert 0 <= posYValue <= 1, f"PosY value should be in range [0, 1], got {posYValue}."
+ assert set(sample.keys()) == expected_keys, (
+ f"{variant}: Expected keys {sorted(expected_keys)}, got {sorted(sample.keys())}."
+ )
- print("All DSprites dataset tests passed successfully!")
+ # Check image type and shape
+ image = sample["image"]
+ assert isinstance(image, Image.Image), f"{variant}: 'image' should be PIL.Image, got {type(image)}."
+ image_np = np.array(image)
+ assert image_np.dtype == np.uint8, f"{variant}: Image dtype should be uint8, got {image_np.dtype}."
+
+ if variant == "original":
+ assert image_np.shape == (64, 64), f"{variant}: Expected shape (64, 64), got {image_np.shape}."
+ else:
+ assert image_np.shape == (64, 64, 3), f"{variant}: Expected shape (64, 64, 3), got {image_np.shape}."
+
+ # Check label structure
+ label = sample["label"]
+ label_values = sample["label_values"]
+ assert isinstance(label, list), f"{variant}: label should be list, got {type(label)}."
+ assert isinstance(label_values, list), f"{variant}: label_values should be list, got {type(label_values)}."
+ assert len(label) == 6, f"{variant}: label should have 6 elements, got {len(label)}."
+ assert len(label_values) == 6, f"{variant}: label_values should have 6 elements, got {len(label_values)}."
+
+ # Check factor ranges
+ assert 0 <= sample["color"] < 1, f"{variant}: color out of range."
+ assert 0 <= sample["shape"] < 3, f"{variant}: shape out of range."
+ assert 0 <= sample["scale"] < 6, f"{variant}: scale out of range."
+ assert 0 <= sample["orientation"] < 40, f"{variant}: orientation out of range."
+ assert 0 <= sample["posX"] < 32, f"{variant}: posX out of range."
+ assert 0 <= sample["posY"] < 32, f"{variant}: posY out of range."
+
+ # Check factor values
+ assert sample["colorValue"] == 1.0, f"{variant}: colorValue should be 1.0."
+ assert sample["shapeValue"] in [1.0, 2.0, 3.0], f"{variant}: shapeValue out of range."
+ assert 0.5 <= sample["scaleValue"] <= 1, f"{variant}: scaleValue out of range."
+ assert 0 <= sample["orientationValue"] <= 2 * np.pi, f"{variant}: orientationValue out of range."
+ assert 0 <= sample["posXValue"] <= 1, f"{variant}: posXValue out of range."
+ assert 0 <= sample["posYValue"] <= 1, f"{variant}: posYValue out of range."
+
+ # Variant-specific checks
+ if variant == "color":
+ color_rgb = sample["colorRGB"]
+ assert isinstance(color_rgb, list), f"{variant}: colorRGB should be list."
+ assert len(color_rgb) == 3, f"{variant}: colorRGB should have 3 elements."
+ for c in color_rgb:
+ assert 0.5 <= c <= 1.0, f"{variant}: colorRGB values should be in [0.5, 1.0], got {c}."
diff --git a/stable_datasets/tests/images/test_dsprites_color.py b/stable_datasets/tests/images/test_dsprites_color.py
deleted file mode 100644
index a911413b..00000000
--- a/stable_datasets/tests/images/test_dsprites_color.py
+++ /dev/null
@@ -1,92 +0,0 @@
-import numpy as np
-from PIL import Image
-
-from stable_datasets.images import DSpritesColor
-
-
-def test_dsprites_dataset():
- # Load training split
- dsprites_train = DSpritesColor(split="train")
-
- # Test 1: Check number of training samples
- expected_num_train_samples = 737280
- assert len(dsprites_train) == expected_num_train_samples, (
- f"Expected {expected_num_train_samples} training samples, got {len(dsprites_train)}."
- )
-
- # Test 2: Check sample keys and label range
- sample = dsprites_train[0]
- expected_keys = {
- "image",
- "index",
- "label",
- "label_values",
- "color",
- "shape",
- "scale",
- "orientation",
- "posX",
- "posY",
- "colorValue",
- "colorRGB",
- "shapeValue",
- "scaleValue",
- "orientationValue",
- "posXValue",
- "posYValue",
- }
- assert sorted(set(sample.keys())) == sorted(set(expected_keys)), (
- f"Expected keys {sorted(set(expected_keys))}, got {sorted(set(sample.keys()))}."
- )
-
- # Test 3: Validate image type
- image = sample["image"]
- assert isinstance(image, Image.Image), f"Image should be a PIL.Image.Image, got {type(image)}."
- image_np = np.array(image)
- assert image_np.ndim == 3, f"DSprites images should be HxW, got shape {image_np.shape}."
- assert image_np.dtype == np.uint8, f"Image dtype should be uint8, got {image_np.dtype}."
- assert image_np.shape == (64, 64, 3), f"Image should have shape (64, 64), got {image_np.shape}"
-
- # Test 4: Validate label type and range
- label = sample["label"]
- label_values = sample["label_values"]
- colorRGB = sample["colorRGB"]
- assert isinstance(label, list), f"Label should be int, got {type(list)}."
- assert isinstance(label_values, list), f"Label values should be list, got {type(label_values)}."
- assert isinstance(colorRGB, list), f"Color RGB should be list, got {type(colorRGB)}."
- assert len(label) == 6, f"Label should have 6 elements, got {len(label)}."
- assert len(label_values) == 6, f"Label values should have 6 elements, got {len(label_values)}."
- assert len(colorRGB) == 3, f"Color RGB should have 3 elements, got {len(colorRGB)}."
-
- color = sample["color"]
- shape = sample["shape"]
- scale = sample["scale"]
- orientation = sample["orientation"]
- posX = sample["posX"]
- posY = sample["posY"]
- colorValue = sample["colorValue"]
- shapeValue = sample["shapeValue"]
- scaleValue = sample["scaleValue"]
- orientationValue = sample["orientationValue"]
- posXValue = sample["posXValue"]
- posYValue = sample["posYValue"]
-
- assert 0 <= color < 1, f"Color should be in range [0, 0], got {color}."
- assert 0 <= shape < 3, f"Shape should be in range [0, 2], got {shape}."
- assert 0 <= scale < 6, f"Scale should be in range [0, 5], got {scale}."
- assert 0 <= orientation < 40, f"Orientation should be in range [0, 39], got {orientation}."
- assert 0 <= posX < 32, f"PosX should be in range [0, 31], got {posX}."
- assert 0 <= posY < 32, f"PosY should be in range [0, 31], got {posY}."
- assert colorValue == 1.0, f"Color value should be 1.0, got {colorValue}."
- assert shapeValue in [1.0, 2.0, 3.0], f"Shape value should be in [1.0, 2.0, 3.0], got {shapeValue}."
- assert 0.5 <= scaleValue <= 1, f"Scale value should be in range [0.5, 1], got {scaleValue}."
- assert 0 <= orientationValue <= 2 * np.pi, (
- f"Orientation value should be in range [0, 2pi], got {orientationValue}."
- )
- assert 0 <= posXValue <= 1, f"PosX value should be in range [0, 1], got {posXValue}."
- assert 0 <= posYValue <= 1, f"PosY value should be in range [0, 1], got {posYValue}."
- assert 0.5 <= colorRGB[0] <= 1.0, f"Color RGB should be in range [0.5, 1.0], got {colorRGB}."
- assert 0.5 <= colorRGB[1] <= 1.0, f"Color RGB should be in range [0.5, 1.0], got {colorRGB}."
- assert 0.5 <= colorRGB[2] <= 1.0, f"Color RGB should be in range [0.5, 1.0], got {colorRGB}."
-
- print("All DSpritesColor dataset tests passed successfully!")
diff --git a/stable_datasets/tests/images/test_dsprites_noise.py b/stable_datasets/tests/images/test_dsprites_noise.py
deleted file mode 100644
index a4e1ce6c..00000000
--- a/stable_datasets/tests/images/test_dsprites_noise.py
+++ /dev/null
@@ -1,85 +0,0 @@
-import numpy as np
-from PIL import Image
-
-from stable_datasets.images import DSpritesNoise
-
-
-def test_dsprites_noise_dataset():
- # Load training split
- dsprites_train = DSpritesNoise(split="train")
-
- # Test 1: Check number of training samples
- expected_num_train_samples = 737280
- assert len(dsprites_train) == expected_num_train_samples, (
- f"Expected {expected_num_train_samples} training samples, got {len(dsprites_train)}."
- )
-
- # Test 2: Check sample keys and label range
- sample = dsprites_train[0]
- expected_keys = {
- "image",
- "index",
- "label",
- "label_values",
- "color",
- "shape",
- "scale",
- "orientation",
- "posX",
- "posY",
- "colorValue",
- "shapeValue",
- "scaleValue",
- "orientationValue",
- "posXValue",
- "posYValue",
- }
- assert sorted(set(sample.keys())) == sorted(set(expected_keys)), (
- f"Expected keys {expected_keys}, got {set(sample.keys())}."
- )
-
- # Test 3: Validate image type
- image = sample["image"]
- assert isinstance(image, Image.Image), f"Image should be a PIL.Image.Image, got {type(image)}."
- image_np = np.array(image)
- assert image_np.ndim == 3, f"DSprites images should be HxWx3, got shape {image_np.shape}."
- assert image_np.dtype == np.uint8, f"Image dtype should be uint8, got {image_np.dtype}."
- assert image_np.shape == (64, 64, 3), f"Image should have shape (64, 64, 3), got {image_np.shape}"
-
- # Test 4: Validate label type and range
- label = sample["label"]
- label_values = sample["label_values"]
- assert isinstance(label, list), f"Label should be list, got {type(label)}."
- assert isinstance(label_values, list), f"Label values should be list, got {type(label_values)}."
- assert len(label) == 6, f"Label should have 6 elements, got {len(label)}."
- assert len(label_values) == 6, f"Label values should have 6 elements, got {len(label_values)}."
-
- color = sample["color"]
- shape = sample["shape"]
- scale = sample["scale"]
- orientation = sample["orientation"]
- posX = sample["posX"]
- posY = sample["posY"]
- colorValue = sample["colorValue"]
- shapeValue = sample["shapeValue"]
- scaleValue = sample["scaleValue"]
- orientationValue = sample["orientationValue"]
- posXValue = sample["posXValue"]
- posYValue = sample["posYValue"]
-
- assert 0 <= color < 1, f"Color should be in range [0, 0], got {color}."
- assert 0 <= shape < 3, f"Shape should be in range [0, 2], got {shape}."
- assert 0 <= scale < 6, f"Scale should be in range [0, 5], got {scale}."
- assert 0 <= orientation < 40, f"Orientation should be in range [0, 39], got {orientation}."
- assert 0 <= posX < 32, f"PosX should be in range [0, 31], got {posX}."
- assert 0 <= posY < 32, f"PosY should be in range [0, 31], got {posY}."
- assert colorValue == 1.0, f"Color value should be 1.0, got {colorValue}."
- assert shapeValue in [1.0, 2.0, 3.0], f"Shape value should be in [1.0, 2.0, 3.0], got {shapeValue}."
- assert 0.5 <= scaleValue <= 1, f"Scale value should be in range [0.5, 1], got {scaleValue}."
- assert 0 <= orientationValue <= 2 * np.pi, (
- f"Orientation value should be in range [0, 2pi], got {orientationValue}."
- )
- assert 0 <= posXValue <= 1, f"PosX value should be in range [0, 1], got {posXValue}."
- assert 0 <= posYValue <= 1, f"PosY value should be in range [0, 1], got {posYValue}."
-
- print("All DSpritesNoise dataset tests passed successfully!")
diff --git a/stable_datasets/tests/images/test_dsprites_scream.py b/stable_datasets/tests/images/test_dsprites_scream.py
deleted file mode 100644
index 6652fc2c..00000000
--- a/stable_datasets/tests/images/test_dsprites_scream.py
+++ /dev/null
@@ -1,85 +0,0 @@
-import numpy as np
-from PIL import Image
-
-from stable_datasets.images import DSpritesScream
-
-
-def test_dsprites_scream_dataset():
- # Load training split
- dsprites_train = DSpritesScream(split="train")
-
- # Test 1: Check number of training samples
- expected_num_train_samples = 737280
- assert len(dsprites_train) == expected_num_train_samples, (
- f"Expected {expected_num_train_samples} training samples, got {len(dsprites_train)}."
- )
-
- # Test 2: Check sample keys and label range
- sample = dsprites_train[0]
- expected_keys = {
- "image",
- "index",
- "label",
- "label_values",
- "color",
- "shape",
- "scale",
- "orientation",
- "posX",
- "posY",
- "colorValue",
- "shapeValue",
- "scaleValue",
- "orientationValue",
- "posXValue",
- "posYValue",
- }
- assert sorted(set(sample.keys())) == sorted(set(expected_keys)), (
- f"Expected keys {expected_keys}, got {set(sample.keys())}."
- )
-
- # Test 3: Validate image type
- image = sample["image"]
- assert isinstance(image, Image.Image), f"Image should be a PIL.Image.Image, got {type(image)}."
- image_np = np.array(image)
- assert image_np.ndim == 3, f"DSprites images should be HxWx3, got shape {image_np.shape}."
- assert image_np.dtype == np.uint8, f"Image dtype should be uint8, got {image_np.dtype}."
- assert image_np.shape == (64, 64, 3), f"Image should have shape (64, 64, 3), got {image_np.shape}"
-
- # Test 4: Validate label type and range
- label = sample["label"]
- label_values = sample["label_values"]
- assert isinstance(label, list), f"Label should be list, got {type(label)}."
- assert isinstance(label_values, list), f"Label values should be list, got {type(label_values)}."
- assert len(label) == 6, f"Label should have 6 elements, got {len(label)}."
- assert len(label_values) == 6, f"Label values should have 6 elements, got {len(label_values)}."
-
- color = sample["color"]
- shape = sample["shape"]
- scale = sample["scale"]
- orientation = sample["orientation"]
- posX = sample["posX"]
- posY = sample["posY"]
- colorValue = sample["colorValue"]
- shapeValue = sample["shapeValue"]
- scaleValue = sample["scaleValue"]
- orientationValue = sample["orientationValue"]
- posXValue = sample["posXValue"]
- posYValue = sample["posYValue"]
-
- assert 0 <= color < 1, f"Color should be in range [0, 0], got {color}."
- assert 0 <= shape < 3, f"Shape should be in range [0, 2], got {shape}."
- assert 0 <= scale < 6, f"Scale should be in range [0, 5], got {scale}."
- assert 0 <= orientation < 40, f"Orientation should be in range [0, 39], got {orientation}."
- assert 0 <= posX < 32, f"PosX should be in range [0, 31], got {posX}."
- assert 0 <= posY < 32, f"PosY should be in range [0, 31], got {posY}."
- assert colorValue == 1.0, f"Color value should be 1.0, got {colorValue}."
- assert shapeValue in [1.0, 2.0, 3.0], f"Shape value should be in [1.0, 2.0, 3.0], got {shapeValue}."
- assert 0.5 <= scaleValue <= 1, f"Scale value should be in range [0.5, 1], got {scaleValue}."
- assert 0 <= orientationValue <= 2 * np.pi, (
- f"Orientation value should be in range [0, 2pi], got {orientationValue}."
- )
- assert 0 <= posXValue <= 1, f"PosX value should be in range [0, 1], got {posXValue}."
- assert 0 <= posYValue <= 1, f"PosY value should be in range [0, 1], got {posYValue}."
-
- print("All DSpritesScream dataset tests passed successfully!")
diff --git a/stable_datasets/tests/images/test_eurosat.py b/stable_datasets/tests/images/test_eurosat.py
new file mode 100644
index 00000000..cede0807
--- /dev/null
+++ b/stable_datasets/tests/images/test_eurosat.py
@@ -0,0 +1,61 @@
+import numpy as np
+import pytest
+from PIL import Image
+
+from stable_datasets.images import EuroSAT
+
+
+pytestmark = pytest.mark.large
+
+
+_EXPECTED_COUNTS = {
+ "train": 16200,
+ "validation": 5400,
+ "test": 5400,
+}
+
+
+@pytest.mark.parametrize("split", ["train", "validation", "test"])
+def test_eurosat_dataset(split):
+ ds = EuroSAT(split=split)
+
+ # Test 1: Check sample count
+ expected = _EXPECTED_COUNTS[split]
+ assert len(ds) == expected, f"[{split}] expected {expected} samples, got {len(ds)}."
+
+ # Test 2: Check sample keys
+ sample = ds[0]
+ expected_keys = {"image", "label"}
+ assert set(sample.keys()) == expected_keys, f"[{split}] expected keys {expected_keys}, got {set(sample.keys())}."
+
+ # Test 3: Validate image type and dimensions
+ image = sample["image"]
+ assert isinstance(image, Image.Image), f"[{split}] image should be a PIL.Image.Image, got {type(image)}."
+
+ image_np = np.array(image)
+ assert image_np.shape == (64, 64, 3), f"[{split}] EuroSAT images should be 64x64x3, got {image_np.shape}."
+ assert image_np.dtype == np.uint8, f"[{split}] image dtype should be uint8, got {image_np.dtype}."
+
+ # Test 4: Validate label
+ label = sample["label"]
+ assert isinstance(label, int), f"[{split}] label should be int, got {type(label)}."
+ assert 0 <= label < 10, f"[{split}] label should be in [0, 9], got {label}."
+
+ # Test 5: Verify class names
+ expected_class_names = [
+ "AnnualCrop",
+ "Forest",
+ "HerbaceousVegetation",
+ "Highway",
+ "Industrial",
+ "Pasture",
+ "PermanentCrop",
+ "Residential",
+ "River",
+ "SeaLake",
+ ]
+ assert ds.features["label"].names == expected_class_names, (
+ f"[{split}] class name mismatch. Expected {expected_class_names}, got {ds.features['label'].names}."
+ )
+
+ print(f"All EuroSAT[{split}] dataset tests passed successfully!")
diff --git a/stable_datasets/tests/images/test_med_mnist.py b/stable_datasets/tests/images/test_med_mnist.py
index ec962659..caeca052 100644
--- a/stable_datasets/tests/images/test_med_mnist.py
+++ b/stable_datasets/tests/images/test_med_mnist.py
@@ -57,3 +57,55 @@ def test_med_mnist_variants_download_and_format():
ds_test = MedMNIST(split="test", config_name=variant)
assert len(ds_test) > 0, f"{variant} test dataset should not be empty."
+
+
+def test_med_mnist_larger_sizes():
+ """
+ Download/integration test for MedMNIST+ larger size variants.
+ 2D variants use size=224, 3D variant uses size=64.
+ """
+
+ variants_with_size = [
+ ("pathmnist", 64, False),
+ ("chestmnist", 64, True),
+ ("organmnist3d", 64, False),
+ ]
+
+ for variant, size, is_multi_label in variants_with_size:
+ ds = MedMNIST(split="train", config_name=variant, size=size)
+ assert len(ds) > 0, f"{variant} (size={size}) training dataset should not be empty."
+
+ sample = ds[0]
+ assert set(sample.keys()) == {"image", "label"}
+
+ image = sample["image"]
+ label = sample["label"]
+
+ if variant.endswith("3d"):
+ image_np = np.asarray(image)
+ assert image_np.shape == (size, size, size), (
+ f"{variant}: expected image shape {(size,) * 3}, got {image_np.shape}."
+ )
+ else:
+ assert isinstance(image, Image.Image), (
+ f"{variant}: 'image' should be a PIL.Image object, got {type(image)}."
+ )
+ image_np = np.asarray(image)
+ assert image_np.shape[0] == size and image_np.shape[1] == size, (
+ f"{variant}: expected {size}x{size} image, got {image_np.shape}."
+ )
+
+ if is_multi_label:
+ label_np = np.asarray(label)
+ assert label_np.ndim == 1, f"{variant}: expected 1D multi-label, got shape {label_np.shape}."
+ assert set(np.unique(label_np).tolist()).issubset({0, 1}), (
+ f"{variant}: labels must be 0/1, got {set(np.unique(label_np).tolist())}."
+ )
+ else:
+ assert isinstance(label, int | np.integer), f"{variant}: label should be an int, got {type(label)}."
+
+ ds_val = MedMNIST(split="validation", config_name=variant, size=size)
+ assert len(ds_val) > 0, f"{variant} (size={size}) validation dataset should not be empty."
+
+ ds_test = MedMNIST(split="test", config_name=variant, size=size)
+ assert len(ds_test) > 0, f"{variant} (size={size}) test dataset should not be empty."
diff --git a/stable_datasets/tests/images/test_oxford_pet.py b/stable_datasets/tests/images/test_oxford_pet.py
new file mode 100644
index 00000000..1bf7539b
--- /dev/null
+++ b/stable_datasets/tests/images/test_oxford_pet.py
@@ -0,0 +1,48 @@
+import numpy as np
+import pytest
+from PIL import Image
+
+from stable_datasets.images import OxfordPet
+
+
+pytestmark = pytest.mark.large
+
+
+def test_oxford_pet_dataset():
+ # Load training split
+ pet_train = OxfordPet(split="train")
+
+ # Test 1: Check number of training samples
+ expected_num_train_samples = 3680
+ assert len(pet_train) == expected_num_train_samples, (
+ f"Expected {expected_num_train_samples} training samples, got {len(pet_train)}."
+ )
+
+ # Test 2: Check sample keys
+ sample = pet_train[0]
+ expected_keys = {"image", "label"}
+ assert set(sample.keys()) == expected_keys, f"Expected keys {expected_keys}, got {set(sample.keys())}."
+
+ # Test 3: Validate image type
+ image = sample["image"]
+ assert isinstance(image, Image.Image), f"Image should be a PIL.Image.Image, got {type(image)}."
+
+ # Convert to numpy for basic sanity checks
+ image_np = np.array(image)
+ assert image_np.ndim == 3, f"OxfordPet images should be HxWxC, got shape {image_np.shape}."
+ assert image_np.shape[2] == 3, f"OxfordPet images should have 3 channels, got {image_np.shape[2]}."
+ assert image_np.dtype == np.uint8, f"Image dtype should be uint8, got {image_np.dtype}."
+
+ # Test 4: Validate label type and range
+ label = sample["label"]
+ assert isinstance(label, int), f"Label should be int, got {type(label)}."
+ assert 0 <= label < 37, f"Label should be in range [0, 36], got {label}."
+
+ # Test 5: Load and validate test split
+ pet_test = OxfordPet(split="test")
+ expected_num_test_samples = 3669
+ assert len(pet_test) == expected_num_test_samples, (
+ f"Expected {expected_num_test_samples} test samples, got {len(pet_test)}."
+ )
+
+ print("All OxfordPet dataset tests passed successfully!")
diff --git a/stable_datasets/tests/images/test_plant_village.py b/stable_datasets/tests/images/test_plant_village.py
new file mode 100644
index 00000000..aed19527
--- /dev/null
+++ b/stable_datasets/tests/images/test_plant_village.py
@@ -0,0 +1,55 @@
+import numpy as np
+import pytest
+from PIL import Image
+
+from stable_datasets.images import PlantVillage
+
+
+pytestmark = pytest.mark.large
+
+
+_EXPECTED_COUNTS = {
+ "color": {"train": 43596, "test": 10709},
+ "grayscale": {"train": 43203, "test": 11102},
+ "segmented": {"train": 42984, "test": 11322},
+}
+
+
+@pytest.mark.parametrize("variant", ["color", "grayscale", "segmented"])
+def test_plant_village_dataset(variant):
+ # Test 1: Load training split and check count
+ pv_train = PlantVillage(split="train", config_name=variant)
+ expected_train = _EXPECTED_COUNTS[variant]["train"]
+ assert len(pv_train) == expected_train, (
+ f"[{variant}] expected {expected_train} training samples, got {len(pv_train)}."
+ )
+
+ # Test 2: Check sample keys
+ sample = pv_train[0]
+ expected_keys = {"image", "label", "crop", "disease"}
+ assert set(sample.keys()) == expected_keys, f"[{variant}] expected keys {expected_keys}, got {set(sample.keys())}."
+
+ # Test 3: Validate image type
+ image = sample["image"]
+ assert isinstance(image, Image.Image), f"[{variant}] image should be a PIL.Image.Image, got {type(image)}."
+ image_np = np.array(image)
+ assert image_np.ndim == 3, f"[{variant}] images should be HxWxC, got shape {image_np.shape}."
+ assert image_np.shape[2] == 3, f"[{variant}] images should have 3 channels, got {image_np.shape[2]}."
+ assert image_np.dtype == np.uint8, f"[{variant}] image dtype should be uint8, got {image_np.dtype}."
+
+ # Test 4: Validate label/crop/disease
+ label = sample["label"]
+ assert isinstance(label, int), f"[{variant}] label should be int, got {type(label)}."
+ assert 0 <= label < 38, f"[{variant}] label should be in [0, 37], got {label}."
+
+ crop = sample["crop"]
+ disease = sample["disease"]
+ assert isinstance(crop, str) and crop, f"[{variant}] crop must be a non-empty string, got {crop!r}."
+ assert isinstance(disease, str) and disease, f"[{variant}] disease must be a non-empty string, got {disease!r}."
+
+ # Test 5: Load test split and check count
+ pv_test = PlantVillage(split="test", config_name=variant)
+ expected_test = _EXPECTED_COUNTS[variant]["test"]
+ assert len(pv_test) == expected_test, f"[{variant}] expected {expected_test} test samples, got {len(pv_test)}."
+
+ print(f"All PlantVillage[{variant}] dataset tests passed successfully!")
diff --git a/stable_datasets/tests/images/test_stanford_dogs.py b/stable_datasets/tests/images/test_stanford_dogs.py
new file mode 100644
index 00000000..dccca70e
--- /dev/null
+++ b/stable_datasets/tests/images/test_stanford_dogs.py
@@ -0,0 +1,45 @@
+import numpy as np
+import pytest
+from PIL import Image
+
+from stable_datasets.images import StanfordDogs
+
+
+pytestmark = pytest.mark.large
+
+
+def test_stanford_dogs_dataset():
+ # Test 1: Train split count
+ sd_train = StanfordDogs(split="train")
+ expected_train = 12000
+ assert len(sd_train) == expected_train, f"Expected {expected_train} training samples, got {len(sd_train)}."
+
+ # Test 2: Sample keys
+ sample = sd_train[0]
+ expected_keys = {"image", "label"}
+ assert set(sample.keys()) == expected_keys, f"Expected keys {expected_keys}, got {set(sample.keys())}."
+
+ # Test 3: Image type/shape
+ image = sample["image"]
+ assert isinstance(image, Image.Image), f"Image should be a PIL.Image.Image, got {type(image)}."
+ image_np = np.array(image)
+ assert image_np.ndim == 3, f"StanfordDogs images should be HxWxC, got shape {image_np.shape}."
+ assert image_np.shape[2] == 3, f"StanfordDogs images should have 3 channels, got {image_np.shape[2]}."
+ assert image_np.dtype == np.uint8, f"Image dtype should be uint8, got {image_np.dtype}."
+
+ # Test 4: Label
+ label = sample["label"]
+ assert isinstance(label, int), f"Label should be int, got {type(label)}."
+ assert 0 <= label < 120, f"Label should be in range [0, 119], got {label}."
+
+ # Test 5: Test split count
+ sd_test = StanfordDogs(split="test")
+ expected_test = 8580
+ assert len(sd_test) == expected_test, f"Expected {expected_test} test samples, got {len(sd_test)}."
+
+ # Test 6: Class label features
+ assert len(sd_train.features["label"].names) == 120, (
+ f"Expected 120 breed names, got {len(sd_train.features['label'].names)}."
+ )
+
+ print("All StanfordDogs dataset tests passed successfully!")
diff --git a/stable_datasets/utils.py b/stable_datasets/utils.py
index 025e138b..1e3717bf 100644
--- a/stable_datasets/utils.py
+++ b/stable_datasets/utils.py
@@ -1,3 +1,4 @@
+import copy
import hashlib
import multiprocessing
import os
@@ -143,7 +144,10 @@ def _wrapped_source(self):
cls._source = _wrapped_source # type: ignore[method-assign]
def __init__(self, config_name: str | None = None, **kwargs):
- """Initialize builder, selecting a BuilderConfig if applicable."""
+ """Initialize builder, selecting a BuilderConfig and applying per-instance overrides.
+
+ Any extra ``kwargs`` are applied as overrides on a shallow copy of the selected BuilderConfig.
+ """
if self.BUILDER_CONFIGS:
if config_name is None:
config_name = self.DEFAULT_CONFIG_NAME or self.BUILDER_CONFIGS[0].name
@@ -151,8 +155,24 @@ def __init__(self, config_name: str | None = None, **kwargs):
if not matched:
available = [c.name for c in self.BUILDER_CONFIGS]
raise ValueError(f"Unknown config '{config_name}'. Available: {available}")
- self.config = matched[0]
+
+ if kwargs:
+ config = copy.copy(matched[0])
+ for k, v in kwargs.items():
+ if not hasattr(config, k):
+ valid = sorted(name for name in vars(config) if not name.startswith("_"))
+ raise ValueError(
+ f"Unknown config option {k!r} for config {config_name!r}; valid fields: {valid}"
+ )
+ setattr(config, k, v)
+ self.config = config
+ else:
+ self.config = matched[0]
else:
+ if kwargs:
+ raise ValueError(
+ f"{type(self).__name__} does not declare BUILDER_CONFIGS; got unexpected kwargs: {list(kwargs)}"
+ )
self.config = BuilderConfig(name=config_name or "default")
# Populate self.info / self._dataset_info so _generate_examples can use it
@@ -282,12 +302,29 @@ def __new__(cls, *args, split=None, processed_cache_dir=None, download_dir=None,
features = instance._dataset_info.features
splits_data = {}
+ # Encode config overrides into the cache fingerprint.
+ template = next(
+ (c for c in cls.BUILDER_CONFIGS if c.name == instance.config.name),
+ None,
+ )
+ if template is not None and template is not instance.config:
+ _missing = object()
+ diffs = {
+ k: v
+ for k, v in vars(instance.config).items()
+ if not k.startswith("_") and getattr(template, k, _missing) != v
+ }
+ extra_fp = ":".join(f"{k}={v!r}" for k, v in sorted(diffs.items()))
+ else:
+ extra_fp = ""
+
for sg in split_generators:
shard_dir_name = cache_fingerprint(
cls.__name__,
str(cls.VERSION),
instance.config.name,
sg.name,
+ extra=extra_fp,
)
shard_dir = instance._processed_cache_dir / shard_dir_name